Budgetting bridges

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

using Z = long long int;

struct Edge {
    Z v;
    Z w;
    Z c;
};

Z n, m, k;
vector<Edge> edges;

vector<bool> ret;

struct Level {
    Z c;
    vector<pair<Z, Z>> edges;
    vector<pair<Z, pair<Z, Z>>> plans;
};
vector<Level> levels;

struct UnionFind {
    UnionFind(Z n) : parent(n) {
        for(Z i = 0; i < n; ++i) {
            parent[i] = i;
        }
    }

    Z find(Z x) {
        if(parent[x] != x) {
            log.push_back({x, parent[x]});
            parent[x] = find(parent[x]);
        }
        return parent[x];
    }

    void merge(Z x, Z y) {
        x = find(x);
        y = find(y);
        if(parent[x] != y) {
            log.push_back({x, parent[x]});
            parent[x] = y;
        }
    }

    void reset(Z logPos) {
        if(logPos > (Z)log.size()) {
            throw 5;
        }

        for(Z i = (Z)log.size() - 1; i >= logPos; --i) {
            parent[log[i].first] = log[i].second;
        }
        log.resize(logPos);
    }

    vector<Z> parent;
    vector<pair<Z, Z>> log;
};

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

    cin >> n >> m >> k;

    edges.resize(m);
    map<pair<Z, Z>, Z> lookup;
    for(Z i = 0; i < m; ++i) {
        Edge edge;
        cin >> edge.v >> edge.w >> edge.c;
        --edge.v;
        --edge.w;
        lookup[{edge.v, edge.w}] = edge.c;
        edges.push_back(edge);
    }

    ret.resize(n, true);

    std::sort(edges.begin(), edges.end(), [&](const Edge& a, const Edge& b) { return a.c < b.c; });
    for(const Edge& edge : edges) {
        if(levels.empty() || levels.back().c != edge.c) {
            levels.push_back({edge.c});
        }
        levels.back().edges.push_back({edge.v, edge.w});
    }

    for(Z i = 0; i < k; ++i) {
        Z b;
        cin >> b;

        for(Z j = 0; j < b; ++j) {
            Z v, w;
            cin >> v >> w;
            --v;
            --w;
            auto it = lookup.find({v, w});
            if(it == lookup.end()) throw 5;
            Z c = it->second;

            Z A = 0;
            Z B = (Z)levels.size() - 1;
            while(A != B) {
                Z M = (A + B) / 2;
                if(levels[M].c < c) {
                    A = M + 1;
                } else {
                    B = M;
                }
            }
            if(levels[A].c != c) throw 5;
            levels[A].plans.push_back({i, {v, w}});
        }
    }

    UnionFind uf(n);
    for(Level& level : levels) {
        std::sort(level.plans.begin(), level.plans.end());

        Z logPos = (Z)uf.log.size();
        Z ci = -1;
        for(auto elem : level.plans) {
            if(elem.first != ci) {
                uf.reset(logPos);
            }
            ci = elem.first;
            auto [v, w] = elem.second;
            if(uf.find(v) == uf.find(w)) {
                ret[ci] = false;
            }
            uf.merge(v, w);
        }
        uf.reset(logPos);

        for(auto p : level.edges) {
            uf.merge(p.first, p.second);
        }
    }

    for(Z i = 0; i < k; ++i) {
        cout << (ret[i] ? "YES" : "NO") << "\n";
    }

    return 0;
}

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