問題概要

$N$ 個の都市があって、都市間を $K$ 本の「道路」と $L$ 本の「鉄道」が結んでいる。各道路と各鉄道は、結んでいる都市間を双方向に移動することができる。

各都市 $i$ に対して、以下の条件を満たす都市の個数を求めてください。

都市 $i$ から出発して、道路のみを用いてその都市に到達できる
都市 $i$ から出発して、鉄道のみを用いてその都市に到達できる

制約

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

まず「道路」のみの問題を解く

グラフの問題だが、今回は各頂点を結ぶ辺として「道路」と「鉄道」という二種類のものがあるから難しい。

こういうときは、まずは単純な問題から考えるのがいい。つまり、道路だけの問題を考えてみよう。まず、道路で繋がれている 2 つの都市を同じグループにするように、Union-Find のマージ操作をしていこう。このとき、次のことが言える。

都市 $a, b$ が鉄道のみで行き来できる ⇔ uf.root(a) == uf.root(b) である

よって、各都市を「Union-Find 上のグループの根」で分類することで問題が解けることがわかった。たとえば $N = 7$ で、都市 0, 1, 2, 3, 4, 5, 6 の根がそれぞれ 5, 1, 5, 5, 4, 5, 1 であった場合には

根が 5 の都市：(0, 2, 3, 5)
根が 1 の都市：(1, 6)
根が 4 の都市：(4)

というように各都市を分類できる。そうすれば、各都市ごとに「根」を求めて、その根に付随する都市の個数を答えればよいことがわかる。

「道路」と「鉄道」

ここまで分かれば、「鉄道」が加わっても同様に解くことができる。まず

「道路」で結ばれた都市グループを管理する Union-Find (uf1 とする)
「鉄道」で結ばれた都市グループを管理する Union-Find (uf2 とする)

という、2 つの Union-Find を用意しよう (ここの発想が少し難しい)。このとき、次が成り立つ。

都市 $a, b$ が「道路のみ」でも「鉄道のみ」でも行き来できる
⇔ uf1.root(a) == uf1.root(b) かつ uf2.root(a) == uf2.root(b)

よって、各都市を

(道路に関する Union-Find 上での根の番号, 鉄道に関する Union-Find 上での根の番号)

というペア値によって分類することで解ける。計算量は、上に述べた「各都市をペア値に分類」という部分を std::map によって実現した場合、 $O(N \log N + (K + L)\alpha(N))$ となる。

コード

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

// Union-Find
struct UnionFind {
    vector<int> par;

    UnionFind() { }
    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() {
    // 入力と UnionFind
    int N, K, L;
    cin >> N >> K >> L;
    UnionFind road(N), train(N);
    for (int i = 0; i < K; ++i) {
        int p, q;
        cin >> p >> q;
        --p, --q;
        road.merge(p, q);
    }
    for (int i = 0; i < L; ++i) {
        int r, s;
        cin >> r >> s;
        --r, --s;
        train.merge(r, s);
    }

    // 各都市を分類する
    map<pair<int,int>,int> nums;
    for (int i = 0; i < N; ++i) {
        int root_road = road.root(i);
        int root_train = train.root(i);
        nums[make_pair(root_road, root_train)]++;
    }

    // 答え
    for (int i = 0; i < N; ++i) {
        int root_road = road.root(i);
        int root_train = train.root(i);
        cout << nums[make_pair(root_road, root_train)] << " ";
    }
    cout << endl;
}