Submission #1255753

#	Time	Username	Problem	Language	Result	Execution time	Memory
1255753		AnhPham	Toll (APIO13_toll)	C++20	0 / 100	0 ms	320 KiB

toll

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll INF = (ll)1e15;

struct Edge{int u,v; ll w; int id; bool isNew;};

int N,M,K;
vector<Edge> oldEdges;
vector<pair<int,int>> newEdges;
vector<ll> people;

struct DSU{vector<int> p; DSU(int n=0){p.resize(n+1); iota(p.begin(), p.end(), 0);} int find(int x){return p[x]==x?x:p[x]=find(p[x]);} bool unite(int a,int b){a=find(a); b=find(b); if(a==b) return false; p[a]=b; return true;}};

// Build MST given mask: chosen new edges have weight 0, unchosen weight INF
// Return whether all chosen edges are in MST, and the MST adjacency list

bool buildMST(int mask, vector<vector<pair<int,ll>>>& g, vector<int>& newEdgeInMST){
    vector<Edge> edges;
    for(auto &e: oldEdges) edges.push_back(e);
    for(int i=0;i<K;i++){
        ll w = (mask>>i & 1)? 0 : INF;
        edges.push_back({newEdges[i].first, newEdges[i].second, w, i, true});
    }
    // sort by weight, tie-break: prefer old edges (isNew false)
    sort(edges.begin(), edges.end(), [](const Edge &a, const Edge &b){ if(a.w!=b.w) return a.w < b.w; return a.isNew < b.isNew; });
    DSU dsu(N);
    g.assign(N+1, {});
    newEdgeInMST.assign(K, 0);
    int cnt=0;
    for(auto &e: edges){
        if(dsu.unite(e.u, e.v)){
            g[e.u].push_back({e.v, e.w});
            g[e.v].push_back({e.u, e.w});
            if(e.isNew) newEdgeInMST[e.id]=1;
            cnt++; if(cnt==N-1) break;
        }
    }
    if(cnt!=N-1) return false;
    // check all chosen present
    for(int i=0;i<K;i++) if(mask>>i &1){ if(!newEdgeInMST[i]) return false; }
    return true;
}

// LCA prep to get max old-edge weight on path between u and v
int LOG;
vector<vector<int>> up;
vector<vector<ll>> upMaxOld; // max old-edge weight on path to ancestor
vector<int> depth;

void dfsPrep(int u,int p, vector<vector<pair<int,ll>>>& g){
    for(auto [v,w]: g[u]) if(v!=p){
        depth[v]=depth[u]+1;
        up[0][v]=u;
        // if edge is from oldEdges, its weight < INF; new edges chosen have weight 0
        // we only want to consider old edges' weights for price cap, so store w if w<INF else 0
        upMaxOld[0][v] = (w>=INF ? 0 : w);
        dfsPrep(v,u,g);
    }
}

ll maxOldOnPath(int u,int v){
    if(depth[u]<depth[v]) swap(u,v);
    ll ans=0;
    int diff = depth[u]-depth[v];
    for(int i=0;i<LOG;i++) if(diff>>i &1){ ans = max(ans, upMaxOld[i][u]); u = up[i][u]; }
    if(u==v) return ans;
    for(int i=LOG-1;i>=0;i--) if(up[i][u]!=up[i][v]){
        ans = max(ans, upMaxOld[i][u]); ans = max(ans, upMaxOld[i][v]);
        u = up[i][u]; v = up[i][v];
    }
    ans = max(ans, upMaxOld[0][u]); ans = max(ans, upMaxOld[0][v]);
    return ans;
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin>>N>>M>>K;
    oldEdges.resize(M);
    for(int i=0;i<M;i++){int u,v; ll w; cin>>u>>v>>w; oldEdges[i]={u,v,w,i,false};}
    newEdges.resize(K);
    for(int i=0;i<K;i++) cin>>newEdges[i].first>>newEdges[i].second;
    people.assign(N+1,0);
    for(int i=1;i<=N;i++) cin>>people[i];

    ll best=0;
    int subsets = 1<<K;
    vector<vector<pair<int,ll>>> g; vector<int> newInMST;

    for(int mask=0; mask<subsets; mask++){
        if(!buildMST(mask, g, newInMST)) continue;
        // prepare LCA
        LOG = 1; while((1<<LOG) <= N) LOG++;
        up.assign(LOG, vector<int>(N+1, 0));
        upMaxOld.assign(LOG, vector<ll>(N+1, 0));
        depth.assign(N+1, 0);
        dfsPrep(1,0,g);
        for(int k=1;k<LOG;k++){
            for(int v=1;v<=N;v++){
                up[k][v] = up[k-1][ up[k-1][v] ];
                upMaxOld[k][v] = max(upMaxOld[k-1][v], upMaxOld[k-1][ up[k-1][v] ]);
            }
        }
        // compute subtree sums to get pass counts for edges in MST
        vector<ll> subs(N+1,0);
        function<void(int,int)> dfsSum = [&](int u,int p){
            subs[u]=people[u];
            for(auto [v,w]: g[u]) if(v!=p){ dfsSum(v,u); subs[u]+=subs[v]; }
        };
        dfsSum(1,0);
        ll totalPeople=0; for(int i=1;i<=N;i++) totalPeople+=people[i];
        // For each new edge i that is in MST (newInMST[i]==1 and mask has it chosen), find its price cap
        ll revenue=0;
        // Need to locate the edge in MST corresponding to the new edge: endpoints are newEdges[i]
        for(int i=0;i<K;i++) if( (mask>>i &1) && newInMST[i]){
            int u=newEdges[i].first, v=newEdges[i].second;
            // price cap = max old-edge weight on path u-v in this MST
            ll price = maxOldOnPath(u,v);
            // find passCount: in MST edge corresponding to this new edge, which child side size?
            // We need to find the edge in MST that connects the two parts: that edge may be one where up relation exists.
            // We can determine pass count by walking from deeper node to its parent until reaching LCA, but easier: if we temporarily remove this new edge in MST,
            // the cut corresponds to subtree of one endpoint when treating the MST as rooted at 1 and considering parent relations.
            // We can find LCA and decide which node is deeper when the edge in tree is between parent-child.
            // Let's find LCA
            int a=u, b=v;
            if(depth[a]<depth[b]) swap(a,b);
            // lift a to depth b
            int diff = depth[a]-depth[b];
            for(int t=0;t<LOG;t++) if(diff>>t &1) a = up[t][a];
            while(a!=b){
                int ta=a, tb=b;
                for(int t=LOG-1;t>=0;t--) if(up[t][a]!=up[t][b]){ a=up[t][a]; b=up[t][b]; }
                a = up[0][a]; b = up[0][b];
            }
            int lca = a;
            // Now determine which node is child in tree for the MST-edge corresponding to the new edge.
            // In our MST built, the new edge itself is an edge between endpoints; it appears in adjacency. We need to find which endpoint is deeper when considering parent from dfsPrep.
            int node=-1;
            if(up[0][ newEdges[i].first ] == newEdges[i].second) node = newEdges[i].first;
            else if(up[0][ newEdges[i].second ] == newEdges[i].first) node = newEdges[i].second;
            else {
                // If neither is direct parent of the other (because our root might be elsewhere), we can use depth: the deeper endpoint's subtree when removing edge
                node = (depth[newEdges[i].first] > depth[newEdges[i].second]) ? newEdges[i].first : newEdges[i].second;
            }
            ll side = subs[node];
            ll pass = min(side, totalPeople - side);
            revenue += pass * price;
        }
        best = max(best, revenue);
    }
    cout<<best<<"\n";
    return 0;
}

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