Submission #1231015

#	Time	Username	Problem	Language	Result	Execution time	Memory
1231015		arshjeet	Toll (APIO13_toll)	C++20	0 / 100	3 ms	5952 KiB

toll

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define pb push_back
#define fast() ios_base::sync_with_stdio(false); cin.tie(NULL)

const int N = 1e5+5;
vector<pair<int,int>> g[N];        // old roads: (to, weight)
vector<pair<int,int>> forest[N];   // we’ll force all new roads in here first
int ppl[N], parentDSU[N], szDSU[N], level[N];
int up[18][N], maxE[18][N];         // binary‐lifting for max‐edge queries
int n, m, k;

struct Edge { int u,v,w, idx; };
vector<Edge> oldE, newE;

int findDSU(int x){
    return parentDSU[x]==x?x:parentDSU[x]=findDSU(parentDSU[x]);
}
void uniteDSU(int a,int b){
    a=findDSU(a); b=findDSU(b);
    if(a==b) return;
    if(szDSU[a]<szDSU[b]) swap(a,b);
    parentDSU[b]=a;
    szDSU[a]+=szDSU[b];
}

// 1) Build the forest forcing all new roads in
void build_initial_forest(){
    // reset DSU
    for(int i=1;i<=n;i++){
        parentDSU[i]=i; szDSU[i]=1;
        forest[i].clear();
    }
    // add all new roads with w = -1 so they float to top
    for(auto &e: newE){
        forest[e.u].pb({e.v, -1});
        forest[e.v].pb({e.u, -1});
        uniteDSU(e.u, e.v);
    }
    // then add old roads by increasing weight
    sort(oldE.begin(), oldE.end(), [](auto &a,auto &b){ return a.w < b.w; });
    for(auto &e: oldE){
        if(findDSU(e.u)!=findDSU(e.v)){
            uniteDSU(e.u,e.v);
            forest[e.u].pb({e.v,e.w});
            forest[e.v].pb({e.u,e.w});
        }
    }
}

// 2) DFS to compute subtree‐flows on this forest
int flow[N];
int dfs_flow(int u,int p){
    int sum = ppl[u];
    for(auto &ed: forest[u]){
        int v = ed.first;
        if(v==p) continue;
        sum += dfs_flow(v,u);
    }
    // record for new‐edge endpoints
    flow[u] = sum;
    return sum;
}

// 3) Lift preprocessing on the forest (treat it as a tree now)
void dfs_lca(int u,int p){
    for(auto &ed: forest[u]){
        int v = ed.first, w = ed.second;
        if(v==p) continue;
        level[v] = level[u] + 1;
        up[0][v] = u;
        maxE[0][v] = w;
        dfs_lca(v,u);
    }
}
int query_max_on_path(int a,int b){
    if(level[a]<level[b]) swap(a,b);
    int mx = 0;
    // lift a up
    int diff = level[a]-level[b];
    for(int i=0;i<18;i++){
        if(diff & (1<<i)){
            mx = max(mx, maxE[i][a]);
            a = up[i][a];
        }
    }
    if(a==b) return mx;
    for(int i=17;i>=0;i--){
        if(up[i][a]!=up[i][b]){
            mx = max(mx, maxE[i][a]);
            mx = max(mx, maxE[i][b]);
            a = up[i][a];
            b = up[i][b];
        }
    }
    mx = max(mx, maxE[0][a]);
    mx = max(mx, maxE[0][b]);
    return mx;
}

signed main(){
    fast();
    cin>>n>>m>>k;
    oldE.reserve(m);
    newE.reserve(k);
    for(int i=0;i<m;i++){
        int u,v,w; cin>>u>>v>>w;
        oldE.push_back({u,v,w,-1});
    }
    for(int i=0;i<k;i++){
        int u,v; cin>>u>>v;
        newE.push_back({u,v,0,i});  // we'll fill w later
    }
    for(int i=1;i<=n;i++) cin>>ppl[i];

    // 1. build forest forcing new roads
    build_initial_forest();

    // 2. compute flows
    dfs_flow(1,0);
    // extract each new-road's flow = flow at deeper endpoint
    vector<pair<int,int>> order; // (flow, idx)
    for(auto &e: newE){
        // pick the endpoint with larger depth in the forest
        // but we don't yet have depth ⇒ reuse flow on v or u?
        // Actually, each new road bridges two subtrees; the flow through edge = dfs on one side
        // So we rerun a small DFS per new edge:
        // (but k≤20, so OK)
        // I'll choose v-side:
        memset(flow,0,sizeof flow);
        dfs_flow(1,0);
        // but simpler: record flows on both and pick min:
        int f1 = 0; // sum in subtree if we cut u-v
        // temporarily remove edge u-v:
        // just run a small DFS:
        function<int(int,int,int)> go = [&](int u,int p,int ban_u){
            int s = ppl[u];
            for(auto &ed: forest[u]){
                int v = ed.first;
                if((u==e.u && v==e.v) || (u==e.v && v==e.u)) continue;
                if(v==p) continue;
                s += go(v,u,ban_u);
            }
            return s;
        };
        f1 = go(e.u,0,e.u);
        int f = min(f1, (int)accumulate(ppl+1,ppl+1+n,0LL)-f1);
        order.push_back({f, e.idx});
        flow[e.idx] = f;
    }

    // 3. sort new roads by descending flow
    sort(order.rbegin(), order.rend());

    // prepare for LCA/max queries
    // reset up/maxE
    for(int i=0;i<18;i++)
        for(int v=1;v<=n;v++)
            up[i][v]=maxE[i][v]=0;
    level[1]=0;
    dfs_lca(1,0);
    for(int i=1;i<18;i++){
        for(int v=1;v<=n;v++){
            up[i][v] = up[i-1][ up[i-1][v] ];
            maxE[i][v] = max(maxE[i-1][v], maxE[i-1][ up[i-1][v] ]);
        }
    }

    // 4. fresh DSU, assign tolls
    for(int i=1;i<=n;i++){
        parentDSU[i]=i; szDSU[i]=1;
    }
    int64_t ans = 0;
    for(auto &pr: order){
        int f = pr.first;
        int idx = pr.second;
        auto &e = newE[idx];
        // find min old-edge on path u-v in current DSU-forest
        int cap = query_max_on_path(e.u, e.v);
        ans += f * cap;
        uniteDSU(e.u, e.v);
    }

    cout<<ans<<"\n";
    return 0;
}

• Subtask 1: 0 / 16
• Subtask 2: 0 / 18
• Subtask 3: 0 / 22
• Subtask 4: 0 / 22
• Subtask 5: 0 / 22