Day 004-Q4 | Union-Find

問題

グループ統合と連結判定。N人と Q個のクエリ:

1 u v: u と v を同じグループに
2 u v: u と v が同じグループか出力（Yes/No）

入力形式

N Q
q_type u v
...

制約

$1 \le N \le 10^5$

$1 \le Q \le 10^5$

0-indexed

入出力例

入力例 1

出力例 1

Yes
No
Yes

ヒント (段階的開示)

ヒント1: 方向性

Union-Find（素集合データ構造）を実装。O(α(N)) でほぼ O(1)。

ヒント2: アプローチ

parent[i]=親、find(x)=根（経路圧縮）、union(x,y)=根を統合（ランク統合）。

ヒント3: 誘導

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx == ry: return
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1

模範解答 (Python)

import sys
input = sys.stdin.readline

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1

    def same(self, x, y):
        return self.find(x) == self.find(y)

def solve():
    N, Q = map(int, input().split())
    uf = UnionFind(N)

    results = []
    for _ in range(Q):
        q, u, v = map(int, input().split())
        if q == 1:
            uf.union(u, v)
        else:
            results.append("Yes" if uf.same(u, v) else "No")

    print('\n'.join(results))

solve()

Step-by-Step 解説

1素朴実装
find が親を再帰的にたどるだけだと、木が深くなって O(N) になる。

2経路圧縮
findのたびに parent[x] = find(parent[x]) でツリーを平坦化。

3ランクによる統合
rank が大きい方を根にする → ツリー高さを O(log N) に抑える。

よくあるミス

ミス	原因	正しい書き方
findでキャッシュなし	毎回ルートまで辿る	経路圧縮を実装
unionで親を片寄せ	O(N)になる	ランクで高さ管理
parent比較で判定	別の根の可能性	`find(x) == find(y)`

Day 004-Q4 — Union-Find

問題