Budgetting bridges

#include <bits/stdc++.h>
using namespace std;
struct E {
    int u, v;
    long c;
};
const int N = 1e5+1;
int dsu_p[N], dsu_s[N];
int dsu_get(int x) {
    while (dsu_p[x] != x) x = dsu_p[x];
    return x;
}
struct U {
    int a, b, p, s;
};
U dsu_j(int a, int b) {
    a = dsu_get(a);
    b = dsu_get(b);
    if (a == b) return {0, 0, 0, 0};
    // b -> a
    if (dsu_s[b] > dsu_s[a]) swap(a, b);
    U r{a, b, dsu_p[b], dsu_s[a]};
    dsu_p[b] = a;
    dsu_s[a] += dsu_s[b];
    return r;
}
void do_u(U x) {
    if (x.a == 0) return;
    dsu_p[x.b] = x.p;
    dsu_s[x.a] = x.s;
}
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    int n, m, k;
    cin >> n >> m >> k;
    for (int i = 1; i <= n; ++i) dsu_p[i] = i, dsu_s[i] = 1;
    E e[m];
    int ord[m];
    map<pair<int, int>, int> vtoi;
    for (int i = 0; i < m; ++i) {
        ord[i] = i;
        cin >> e[i].u >> e[i].v >> e[i].c;
        vtoi[{e[i].u, e[i].v}] = i;
    }
    sort(ord, ord+m, [&](int a, int b) {return e[a].c < e[b].c;});
    map<long, map<int, vector<int>>> ctoe;
    bool ans[k];
    for (int i = 0; i < k; ++i) {
        ans[i] = 1;
        int b;
        cin >> b;
        for (int j = 0; j < b; ++j) {
            int u, v;
            cin >> u >> v;
            int ei = vtoi[{u, v}];
            long c = e[ei].c;
            ctoe[c][i].push_back(ei);
        }
    }
    for (int ei = 0; ei < m; ) {
        long cc = e[ord[ei]].c;
        for (const auto& [pi, es] : ctoe[cc]) {
            vector<U> undos;
            for (int pei : es) {
                int a = dsu_get(e[pei].u), b = dsu_get(e[pei].v);
                if (a == b) {
                    ans[pi] = false;
                    break;
                }
                undos.push_back(dsu_j(a, b));
            }
            for (auto& x : undos) do_u(x);
        }
        for (; ei < m && e[ord[ei]].c == cc; ++ei) {
            int a = dsu_get(e[ord[ei]].u), b = dsu_get(e[ord[ei]].v);
            dsu_j(a, b);
        }
    }
    for (int i = 0; i < k; ++i) {
        if (ans[i]) cout << "YES\n";
        else cout << "NO\n";
    }
}

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: ACCEPTED

Test 4

Verdict: ACCEPTED

Test 5

Verdict: ACCEPTED

Test 6

Verdict: ACCEPTED

Test 7

Verdict: ACCEPTED

Test 8

Verdict: ACCEPTED

Test 9

Verdict: ACCEPTED

Test 10

Verdict: ACCEPTED

Test 11

Verdict: ACCEPTED