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] << " "; } }