ABC 355 F

문제의 조건을 살펴보면 간선의 가중치가 10 이하로 매우 작다는 것을 알 수 있다. MST의 기본은 간선을 가중치 기준으로 정렬한 후, Union Find를 돌리는 것이다. 이때, 간선을 정렬할 필요 없이 가중치 별로 저장할 수 있다.

 

Union Find를 10개 만들면 이를 쉽게 관리할 수 있을 것이라 생각했지만 단순히 Union Find를 10개 만드는 것으로는 해결되지 않았다. 여기서 필요했던 또 하나의 아이디어는 Union Find를 누적시키는 것이다. 추가된 간선이 더 낮은 간선들로 연결된 하나의 컴포넌트 내에 있다면, 이를 카운트할 필요가 없다. 따라서, 현재 가중치 이상인 Union Find에 모두 Union을 해주고 현재 가중치 - 1의 Union Find, 현재 가중치의 Union Find에서 컴포넌트의 개수를 관리해주면 해결할 수 있다. 

 

Union Find를 여러 개 쓰는 것 자체는 어렵지 않았지만 Union Find를 누적시킨다는 아이디어가 떠오르지 않아 어렵게 느껴졌다.

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

inline void init() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
}

class union_find {
private:
    vector<ll> p;
public:
    union_find(ll n) {
        p.resize(n + 1);
        iota(p.begin(), p.end(), 0LL);
    }

    ll find(ll x) {
        return x == p[x] ? x : p[x] = find(p[x]);
    }

    bool merge(ll x, ll y) {
        x = find(x), y = find(y);

        if (x == y) {
            return false;
        } else {
            p[y] = x;
            return true;
        }
    }
};

int main() {
    init();

    ll n, q;

    cin >> n >> q;

    vector uf(11, union_find(n));
    vector cnt(11, 0LL);

    for (ll u, v, w, i = 1; i < n; i++) {
        cin >> u >> v >> w;

        for (ll j = w; j <= 10; j++) {
            cnt[j] += uf[j].merge(u, v);
        }
    }

    for (ll u, v, w; q--;) {
        cin >> u >> v >> w;

        for (ll j = w; j <= 10; j++) {
            cnt[j] += uf[j].merge(u, v);
        }

        ll ans = 0;

        for (ll j = 1; j <= 10; j++) {
            ans += j * (cnt[j] - cnt[j - 1]);
        }

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

    return 0;
}