Submission #1274156

#	Time	Username	Problem	Language	Result	Execution time	Memory
1274156		Icelast	Reconstruction Project (JOI22_reconstruction)	C++20	70 / 100	5109 ms	404596 KiB

reconstruction

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

struct DSU {
    int maxn;
    vector<int> pa;
    vector<int> sz;

    DSU(int maxn = 0) { init(maxn); }

    void init(int _maxn) {
        maxn = _maxn;
        if ((int)pa.size() < maxn + 5) {
            pa.resize(maxn + 5);
            sz.resize(maxn + 5);
        }
        reset(maxn);
    }

    // reset to fresh DSU of size n (1..n)
    void reset(int n) {
        // note: loop is necessary (same complexity as original) but avoids reallocations
        for (int i = 0; i <= n; ++i) {
            pa[i] = i;
            sz[i] = 1;
        }
    }

    int find(int x) {
        while (x != pa[x]) {
            pa[x] = pa[pa[x]];
            x = pa[x];
        }
        return x;
    }

    bool merge(int x, int y) {
        x = find(x);
        y = find(y);
        if (x == y) return false;
        if (sz[x] < sz[y]) swap(x, y);
        pa[y] = x;
        sz[x] += sz[y];
        return true;
    }

    bool same(int x, int y) { return find(x) == find(y); }
    int size(int x) { return sz[find(x)]; }
};

struct Edge {
    int u, v;
    int w;
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m;
    if (!(cin >> n >> m)) return 0;

    vector<Edge> e(m + 1);
    for (int i = 1; i <= m; ++i) {
        cin >> e[i].u >> e[i].v >> e[i].w;
    }

    sort(e.begin() + 1, e.end(), [](const Edge &a, const Edge &b) {
        return a.w < b.w;
    });

    // preallocate lt/rt
    vector<vector<int>> lt(m + 1), rt(m + 1);

    DSU dsu(n); // allocate once
    dsu.init(n);

    // build lt (1..m), lt[0] stays empty
    for (int i = 1; i <= m; ++i) {
        dsu.reset(n);
        dsu.merge(e[i].u, e[i].v);

        lt[i].reserve((i > 0 ? lt[i - 1].size() : 0) + 1);
        lt[i].push_back(i);

        const auto &prev = lt[i - 1];
        for (int it : prev) {
            if (dsu.merge(e[it].u, e[it].v)) lt[i].push_back(it);
        }
    }

    // build rt ; rt[m] stays empty in original logic (we mirror original indexing)
    for (int i = m; i >= 1; --i) {
        dsu.reset(n);
        dsu.merge(e[i].u, e[i].v);

        rt[i - 1].reserve((i <= m ? rt[i].size() : 0) + 1);
        rt[i - 1].push_back(i);

        const auto &next = rt[i];
        for (int it : next) {
            if (dsu.merge(e[it].u, e[it].v)) rt[i - 1].push_back(it);
        }
    }

    int Q;
    cin >> Q;
    int p = 0; // number of edges with weight <= x (keeps same meaning as your original code)
    for (int qi = 0; qi < Q; ++qi) {
        int x;
        cin >> x;
        while (p + 1 <= m && e[p + 1].w <= x) ++p;

        dsu.reset(n);
        ll ans = 0;

        const auto &L = lt[p];
        const auto &R = rt[p];
        int p1 = 0, p2 = 0;
        int total = (int)L.size() + (int)R.size();

        // merge edges in the order you had (closest weight to x first)
        for (int t = 0; t < total; ++t) {
            if (p2 == (int)R.size() || (p1 < (int)L.size() && ll(x - e[L[p1]].w) <= ll(e[R[p2]].w - x))) {
                int idx = L[p1++];
                if (dsu.merge(e[idx].u, e[idx].v)) ans += ll(x - e[idx].w);
            } else {
                int idx = R[p2++];
                if (dsu.merge(e[idx].u, e[idx].v)) ans += ll(e[idx].w - x);
            }
        }

        cout << ans << '\n';
    }

    return 0;
}
//chatgpt test

• Subtask 1: 3 / 3
• Subtask 2: 4 / 4
• Subtask 3: 0 / 7
• Subtask 4: 28 / 28
• Subtask 5: 35 / 35
• Subtask 6: 0 / 23