Union-Find（素集合データ構造）とは？仕組み・コード・計算量を解説

Union-Findが解決する問題

ケーブルが1本ずつ追加されていくコンピュータのネットワークを想像してください。ケーブルを追加するたびに、「コンピュータAとBはつながっているか？」に答える必要があります。質問のたびにグラフ探索を実行すると、毎回O(n + m)かかります。Union-Findは、ケーブルの追加に合わせて連結なグループを保持することで、各質問にほぼO(1)で答えます。

同じ要件は、クラスカル法による最小全域木（この2つの端点は、すでに同じ木に属しているか？）、無向グラフの閉路検出、同じメールアドレスを共有するアカウントのグループ化、パーコレーションのシミュレーション、画像処理での連結領域のラベル付けにも現れます。

イメージで理解する

すべての要素は親を指しています。自分自身を指す要素が根で、根がグループ全体の代表です。xのグループを見つけるには、根に着くまで親ポインタをたどります。2つのグループを併合するには、両方の根を見つけ、一方の根をもう一方に向けます。

2つの最適化で、木を平らに保ちます。

• ランク（またはサイズ）による併合：低い（または小さい）木を、高い（または大きい）木の下に付けます。こうすると、同じ大きさの木を併合したときだけ、木の高さが増えます。
• 経路圧縮：findの最中に、経路上のすべてのノードを根に直接向けます。こうすると、次の検索は1ステップで済みます。

具体例：5つの要素

0から4までの要素から始め、それぞれが自分だけのグループに入っています。親の配列は[0, 1, 2, 3, 4]です。サイズによる併合を使い、サイズが等しいときは2つ目の根を子にします。

Pythonでの実装

このバージョンは、サイズによる併合と、経路圧縮付きの反復型findを使うので、再帰の深さの問題がありません。unionは、2つの要素がすでに同じグループにあった場合にFalseを返します。これは閉路検出に便利です。

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n
        self.groups = n

    def find(self, x):
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != root:       # path compression
            self.parent[x], x = root, self.parent[x]
        return root

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False                    # already connected
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra                 # ra is the larger root
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        self.groups -= 1
        return True

    def connected(self, a, b):
        return self.find(a) == self.find(b)

uf = UnionFind(5)
uf.union(0, 1)
uf.union(2, 3)
uf.union(1, 3)
print(uf.find(3), uf.parent)       # 0 [0, 0, 0, 0, 4]
print(uf.connected(1, 4), uf.groups)  # False 2

例：無向グラフの閉路検出

辺を1本ずつ処理します。辺の両端がすでに同じ根を持っているなら、その辺は閉路を作ります。クラスカル法は、辺を重みでソートしたあと、まさにこの判定を使います。

def has_cycle(n, edges):
    uf = UnionFind(n)
    for a, b in edges:
        if not uf.union(a, b):
            return True
    return False

print(has_cycle(4, [(0, 1), (1, 2), (2, 3)]))  # False
print(has_cycle(4, [(0, 1), (1, 2), (2, 0)]))  # True

時間計算量と空間計算量

α(n)は逆アッカーマン関数で、非常にゆっくり増えるため、現実のコンピュータに収まるどんなnでも、せいぜい4です。実用上は、1操作あたり定数時間という意味です。この上界は、1975年にRobert Tarjanが証明しました。空間は、親とサイズの配列でO(n)です。

種類と、よくある間違い

• 根ではなく要素そのものを比較する：unionは、aとbそのものではなく、find(a)とfind(b)をつなぐ必要があります。
• parent[x]が代表だと思い込む。併合を行ったあとで正しいのはfind(x)だけです。親でグループ分けする前に、すべての要素にfindを呼びましょう。
• ランクによる併合なしで、大きな入力に再帰的なfindを使うと、Pythonの再帰の上限に達することがあります。上のような反復版なら避けられます。
• 要素が0からn-1の整数ではない場合：先にdictでインデックスに対応付けるか、親をdictに格納して、要素を遅延的に追加します。
• Union-Findはグループを分割できません。辺が削除される場合は、別のデータ構造か、削除を逆順に処理するオフラインの手法が必要です。
• 重み付きUnion-Findは、各ノードに親へのオフセットを保存し、「aはbよりどれだけ重いか？」といった問いに答えます。

使いどころと、典型的な問題

グループが併合されるだけで、2つの要素が同じグループにあるか、あるいはグループがいくつあるかを知りたいときに、Union-Findを使います。静的なグラフで連結性を1回だけ調べるなら、素のBFSやDFSのほうが単純です。実際の経路が必要なときは、Union-Findが分かるのは所属だけなので、グラフ探索を使いましょう。

典型的な問題：連結成分の数、島の数（グリッドのセルを要素とする）、冗長な接続（閉路を作る辺）、アカウントの併合、a == b, b != cのような等式の充足可能性、クラスカル法による最小全域木などです。