AtCoder ABC 049 D - 連結 / Connectivity (ARC 065 D) (青色, 400 点)

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

問題へのリンク

問題概要

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

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

制約

• 
• 

まず道路のみだったら

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

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

---

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

---

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

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

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

道路と鉄道だったら

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

• 道路に関する Union-Find と、鉄道に関する Union-Find の二つを管理して
• 個の頂点を、(道路に関する 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] << " ";
    }
}