Submission #1151421

#TimeUsernameProblemLanguageResultExecution timeMemory
1151421aintaReconstruction Project (JOI22_reconstruction)C++20
0 / 100
0 ms328 KiB
#include <bits/stdc++.h>
using namespace std;
 
typedef long long ll;
 
// Use a constant INF that is large enough (here 1e18 works well).
const ll INF = 1000000000000000000LL / 1000;  // about 1e18
 
// ---------- 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 variables
int N, M;
 
// We'll 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 want to choose an MST minimizing the total cost sum_{e in T}|w_e - X|.
// The idea is to "merge" the two lists:
//   - Let L = all edges with w <= X (we want to pick them in order of decreasing w, so that X - w is as small as possible).
//   - Let G = all edges with w > X (we want to pick them in order of increasing w, so that w - X is small).
// We have pre–sorted arrays edgesDesc (global) and edgesAsc (global), but they contain all edges.
// Hence, we must “advance” the pointers so that we only consider edges with w <= X from edgesDesc
// and only edges with w > X from edgesAsc.
pair<ll,int> computeMST(ll X) {
    int nEdges = 0;
    ll totCost = 0;
    int cntLE = 0; // count of chosen edges with w <= X
    DSU dsu(N);
    int iAsc = 0, iDesc = 0;
    int szAsc = edgesAsc.size(), szDesc = edgesDesc.size();
 
    while(nEdges < N-1 && (iAsc < szAsc || iDesc < szDesc)) {
        // Advance pointer for edgesAsc until we find one with w > X.
        while(iAsc < szAsc && edgesAsc[iAsc].w <= X) iAsc++;
        // Advance pointer for edgesDesc until we find one with w <= X.
        while(iDesc < szDesc && edgesDesc[iDesc].w > X) iDesc++;
 
        ll costDesc = INF, costAsc = INF;
        if(iDesc < szDesc) {
            // For an edge with w <= X, cost = X - w.
            costDesc = X - edgesDesc[iDesc].w;
        }
        if(iAsc < szAsc) {
            // For an edge with w > X, cost = w - X.
            costAsc = edgesAsc[iAsc].w - X;
        }
 
        if(costDesc == INF && costAsc == INF)
            break;
 
        bool useDesc = (costDesc <= costAsc);
        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++;
        }
    }
    if(nEdges < N-1) return {INF, 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.
    // Include 1 and 10^9 (the endpoints) plus every distinct edge weight.
    set<ll> candSet;
    candSet.insert(1);
    candSet.insert(1000000000LL);
    vector<int> ZZ;
    for(auto &e: edges){
        ZZ.push_back(e.w);
    }
    sort(ZZ.begin(),ZZ.end());
    for(auto &e : edges) {
        candSet.insert(e.w);
    }
    for(int i=0;i<ZZ.size()-1;i++){
        candSet.insert((ZZ[i]+ZZ[i+1])/2);
        candSet.insert((ZZ[i]+ZZ[i+1]+1)/2);
    }

    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 chosen in MST) - (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);
    }
 
    // Process queries.
    int Q; cin >> Q;
    // The Q queries (target widths) are given in strictly increasing order.
    // For each query, binary search to determine the correct candidate interval.
    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...