Union-Find 木に関するすごくいい感じの問題

問題概要

$N$ 個の駅がある。 $M$ 本の道路と $L$ 本の鉄道があって、それぞれ双方向に駅ペアを結んでいる。

各駅に対して、その駅から道路のみを使っていくことができて、かつ鉄道のみを使ってもいくことができる駅の個数を答えよ。

制約

$2 \le N \le 2 \times 10^{5}$
$1 \le M, L \le 10^{5}$

まず道路のみだったら

まずは道路のみだった場合を考えてみよう。道路のつなぎ方に関する連結成分を管理するには Union-Find が適している。

そして、頂点 $v$ と同じグループに属するメンバー数を知りたかったら、

$a, b$ が同じグループである ⇔ uf.root(a) == uf.root(b) である

というのを使えばよい。つまり、

$N$ 個の頂点を根ごとにグループしてあげて
各頂点について属する根を求めて
その根のグループのサイズを求める

という風にしてあげればよい (他にも Union-Find に予めグループサイズを持たせる方法もあります)。

道路と鉄道だったら

ここまで来ればほとんど同じで、今度は

道路に関する Union-Find と、鉄道に関する Union-Find の二つを管理して
$N$ 個の頂点を、(道路に関する Union-Find 上の根, 鉄道に関する Union-Find の根) のペアをキーとしてグループ分けして
各頂点について属するグループのサイズを求める

という風にしてあげれば OK

#include <iostream>
#include <map>
#include <vector>
using namespace std;

struct UnionFind {
    vector<int> par;
    
    UnionFind(int n) : par(n, -1) { }
    void init(int n) { par.assign(n, -1); }
    
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

int main() {
    int N, K, L; cin >> N >> K >> L;
    UnionFind uf1(N), uf2(N);
    for (int i = 0; i < K; ++i) {
        int a, b; cin >> a >> b; --a, --b;
        uf1.merge(a, b);
    }
    for (int i = 0; i < L; ++i) {
        int a, b; cin >> a >> b; --a, --b;
        uf2.merge(a, b);
    }

    using pint = pair<int,int>;
    map<pint, int> ma;
    for (int v = 0; v < N; ++v) {
        pint p(uf1.root(v), uf2.root(v));
        ma[p]++;
    }

    for (int v = 0; v < N; ++v) {
        pint p(uf1.root(v), uf2.root(v));
        cout << ma[p] << " ";
    }
}