Submission #1151416

#TimeUsernameProblemLanguageResultExecution timeMemory
1151416aintaReconstruction Project (JOI22_reconstruction)C++20
0 / 100
0 ms324 KiB
#include <bits/stdc++.h>
using namespace std;
 
typedef long long ll;
 
// ---------- DSU (Disjoint Set Union) for MST ----------
struct DSU {
    vector<int> par;
    DSU(int n) : par(n){
        for (int i = 0; i < n; i++) par[i] = i;
    }
    int find(int a){
        return par[a]==a ? a : par[a]=find(par[a]);
    }
    bool unite(int a, int b){
        a = find(a); b = find(b);
        if(a==b)return false;
        par[b]=a;
        return true;
    }
};
 
// ---------- Edge structure ----------
struct Edge {
    int u, v;
    ll w;
};
 
// Global graph parameters
int N, M;
 
// We store edges in two orders:
//  - edgesAsc: sorted in increasing order of w
//  - edgesDesc: sorted in decreasing order of w
vector<Edge> edgesAsc, edgesDesc;
 
// ---------- MST computation for a fixed target X ----------
// For a given X, the cost to reconstruct an edge of original width w is:
//      cost = (X - w)  if w <= X   (slope +1)
//      cost = (w - X)  if w >  X   (slope -1)
// We “merge” the two sorted lists (one for edges with w <= X and one for w > X)
// in order of increasing |w - X|. Then we run a DSU–based MST algorithm over that merged order.
// We also count how many chosen edges satisfy w <= X.
pair<ll,int> computeMST(ll X) {
    int nEdges = 0;
    ll totCost = 0;
    int cntLE = 0; // count edges with w <= X chosen in MST
    DSU dsu(N);
    int iAsc = 0, iDesc = 0;
    int szAsc = edgesAsc.size(), szDesc = edgesDesc.size();
    // Merge on–the–fly the two lists: from edgesDesc we take those with w <= X (cost = X - w),
    // from edgesAsc we take those with w > X (cost = w - X).
    while(nEdges < N-1 && (iAsc < szAsc || iDesc < szDesc)){
        ll costDesc = LLONG_MAX;
        if(iDesc < szDesc && edgesDesc[iDesc].w <= X)
            costDesc = X - edgesDesc[iDesc].w;
        ll costAsc = LLONG_MAX;
        if(iAsc < szAsc && edgesAsc[iAsc].w > X)
            costAsc = edgesAsc[iAsc].w - X;
        bool useDesc = false;
        if(costDesc <= costAsc) useDesc = true;
        if(costDesc==LLONG_MAX && costAsc==LLONG_MAX)
            break;
 
        Edge cur;
        if(useDesc){
            cur = edgesDesc[iDesc++];
        } else {
            cur = edgesAsc[iAsc++];
        }
 
        if(dsu.unite(cur.u, cur.v)){
            totCost += (cur.w <= X ? (X - cur.w) : (cur.w - X));
            nEdges++;
            if(cur.w <= X) cntLE++;
        }
    }
    // (The graph is connected so we always can choose N-1 edges.)
    if(nEdges < N-1) return {LLONG_MAX,0};
    return {totCost, cntLE};
}
 
// ---------- Main ----------
int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
 
    cin >> N >> M;
    vector<Edge> edges(M);
    for (int i=0; i<M; i++){
        int u,v; ll w;
        cin >> u >> v >> w;
        u--; v--;  // convert to 0-index
        edges[i] = {u,v,w};
    }
    // Build sorted arrays:
    edgesAsc = edges;
    sort(edgesAsc.begin(), edgesAsc.end(), [](const Edge &a, const Edge &b){
        return a.w < b.w;
    });
    edgesDesc = edges;
    sort(edgesDesc.begin(), edgesDesc.end(), [](const Edge &a, const Edge &b){
        return a.w > b.w;
    });
 
    // Build candidate X values.
    // We include 1 and 10^9 (the endpoints) plus every distinct edge weight.
    set<ll> candSet;
    candSet.insert(1);
    candSet.insert(1000000000LL);
    for(auto &e : edges) {
        candSet.insert(e.w);
    }
    vector<ll> candidates(candSet.begin(), candSet.end());
    int K = candidates.size();
    // For each candidate X, compute f(X) = MST cost and the derivative d = 2*(# edges with w<=X) - (N-1)
    vector<ll> fval(K), deriv(K);
    for (int i=0; i<K; i++){
        ll X = candidates[i];
        auto res = computeMST(X);
        fval[i] = res.first;
        deriv[i] = 2LL * res.second - (N - 1);
    }
 
    // Now, between any two candidate points candidates[i] and candidates[i+1],
    // the MST remains the same so that for any X in that interval
    //   f(X) = fval[i] + deriv[i]*(X - candidates[i]).
    // (Because f is convex and piecewise linear.)
 
    // Process queries.
    int Q; cin >> Q;
    // The Q queries (target widths) are given in strictly increasing order.
    // For each query, binary search which candidate interval it falls into.
    for (int qi = 0; qi < Q; qi++){
        ll X; cin >> X;
        int idx = (int)(upper_bound(candidates.begin(), candidates.end(), X) - candidates.begin()) - 1;
        if(idx < 0) idx = 0;
        if(idx >= K) idx = K-1;
        ll ans = fval[idx] + deriv[idx] * (X - candidates[idx]);
        cout << ans << "\n";
    }
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...