AtCoderでUnion-Findをいつ使う?経路圧縮とサイズ併合を実装で理解

読了 約10分 たびすけ
AtCoderのUnion-Findを表す2つの木を統合する図

次に読む記事

関連するテーマの記事を、先に確認できます。

Union-Findを選ぶ条件

Union-Findは、頂点をいくつかの集合に分けた状態を保ち、辺の追加に合わせて集合を併合する問題に向いています。

追加した辺の両端が同じ集合に属するかを繰り返し調べるとき、親配列が保持している集合状態を次の辺でも再利用できます。辺を追加するたびに両端の根を比較し、異なる根だけを併合して連結性を保つため、集合全体を単純な探索で毎回たどり直す代わりに、この構造を選べます。最後に集合の個数や代表元を求める場合にもつながります。

問題文では、次の条件を確認します。

  • 辺を追加しながら、連結性の判定や連結成分数の集計を行う。
  • 必要なのは同じ集合か、集合の代表元、集合の大きさであり、具体的な経路や距離ではない。
  • 基本実装は、辺の追加、集合の併合、連結成分の集約を扱う。任意の辺削除は別手法の対象である。

辺を追加すると、両端の集合を一つにまとめます。いったんまとめた集合を任意の辺削除によって分離する処理は、この基本的なUnion-Findの範囲に含まれません。

親配列と根を小さな例で追う

Union-Findは、各頂点から親をたどって集合の代表に到達する構造を持ちます。親配列 parent の値が自分自身になっている頂点をと呼び、根はその集合の代表元になります。

頂点1から頂点4までを使い、辺を1-2、3-4、1-3の順に追加します。表の頂点番号は1始まりで、parentの各値は「その頂点が直接指す親」を表します。sizeは根に記録した集合の大きさです。

操作 parent(概念上の頂点1〜4) 根のサイズ
初期状態 [1, 2, 3, 4] 根1: 1、根2: 1、根3: 1、根4: 1
辺1-2を追加 [1, 1, 3, 4] 根1: 2、根3: 1、根4: 1
辺3-4を追加 [1, 1, 3, 3] 根1: 2、根3: 2
辺1-3を追加 [1, 1, 1, 3] 根1: 4
概念上のfind(4)(コード上はfind(3))を実行 [1, 1, 1, 1] 根1: 4

表の頂点番号と矢印は、概念上の頂点を示す1始まりです。Pythonの配列とコードの添字は0始まりなので、入力の頂点番号abをそれぞれa - 1b - 1へ変換してからparentを調べます。したがって、表の頂点4に対する概念上のfind(4)は、コード上のfind(3)に対応します。

辺1-3を追加した直後は、表の1始まりの矢印で頂点4から根1までの参照が4→3→1となっています。コードの0始まりでは、同じ参照が3→2→0です。表の最後のfind(4)はコード上のfind(3)としてこの経路をたどり、根0を見つけてparent[3]を0へ書き戻します。

経路圧縮とサイズによる併合

経路圧縮は、根を探す途中で見つけた根を、調べた頂点の親へ直接設定する処理です。完全コードの parent[x] = find(parent[x]) がこの更新に当たり、表で示した4→3→1(コードでは3→2→0)の経路を、次回から4→1(コードでは3→0)へ短くします。

根に到達したときは parent[x] == x が成り立つため、その頂点を返します。根でなければ親について同じ処理を再帰的に行い、返ってきた根を parent[x] に保存します。

サイズによる併合は、二つの根の集合の大きさを比べ、小さい集合の根を大きい集合の根へつなぐ処理です。コードでは if size[a] < size[b] のときに根を入れ替え、その後の parent[b] = a で小さい側を大きい側へ付けます。併合後は size[a] += size[b] で新しい根の大きさを更新します。

二つの根が同じなら、辺を追加しても集合は増えないため、コードは併合を行いません。異なる根だけを併合することで、現在の集合の分かれ方を保ったまま辺を処理できます。

ABC284 Cの問題と公式資料

ABC284 Cは、頂点数N、辺数Mの単純無向グラフについて、連結成分の個数を求める問題です。頂点数Nは1以上100以下、辺数Mは0以上N(N - 1) / 2以下です。

入力の1行目にNとMが与えられ、その後のM行に無向辺の両端の頂点番号が1行ずつ与えられます。出力するのは、すべての辺を考慮した後の連結成分数です。

公式解説には、各連結成分を訪れるDFSやBFSの方法と、辺ごとに集合を併合して最後に代表元を数えるUnion-Findの方法が示されています。公式解説の方針を、親配列の更新として確認できます。

ABC284 Cは、辺を入力順に一つずつ追加し、各辺の両端が同じ集合かを判定しながら、集合の状態を保って最後に連結成分数を数える流れを確認できる問題です。この流れを親配列の更新として追えるため、Union-Findの基本操作を一つの問題で確認できます。

公式サンプル1

5 3
1 2
1 3
4 5

出力は次のとおりです。

2

頂点1、2、3が一つの成分になり、頂点4、5がもう一つの成分になる入力です。

公式サンプル2

5 0

出力は次のとおりです。

5

辺がないため、5頂点がそれぞれ別の成分として数えられます。

公式サンプル3

4 6
1 2
1 3
1 4
2 3
2 4
3 4

出力は次のとおりです。

1

4頂点のすべての組に辺がある完全グラフなので、どの頂点からでも他の頂点へ直接移動できます。そのため、4頂点は一つの連結成分になり、出力は1です。

完全なPython参考コード

次のコードは、入力の読み取り、集合の初期化、辺の併合、最後の連結成分数の出力までを含むPython参考実装です。ACLの公式DSU資料はC++の公式APIを説明したもので、このコードはACLのPython公式実装ではありません。親配列とサイズ配列を使って同じ集合管理の考え方を独立に実装しています。

N, M = map(int, input().split())
parent = list(range(N))
size = [1] * N

def find(x):
    if parent[x] == x:
        return x
    parent[x] = find(parent[x])
    return parent[x]

for _ in range(M):
    a, b = map(int, input().split())
    a, b = find(a - 1), find(b - 1)
    if a != b:
        if size[a] < size[b]:
            a, b = b, a
        parent[b] = a
        size[a] += size[b]

print(len({find(i) for i in range(N)}))

コードの正しさと計算量

最初は各頂点が自分自身を親とするため、頂点ごとに一つの集合ができます。辺を読むたびに find が両端の根を返し、根が異なる場合だけ一方の根を他方へつなぎます。この更新は二つの集合を一つにまとめる操作なので、辺を処理した後も、同じ集合に属する頂点が同じ根へたどり着く状態が保たれます。

根が同じ場合は、その辺の両端がすでに同じ連結成分に属しているため、集合の数は変わりません。根が異なる場合は二つの連結成分を併合するため、連結成分の分かれ方と親配列の集合が一致します。すべての辺を処理した後、集合 {find(i) for i in range(N)} には各連結成分の根が一つずつ入るので、その要素数が答えになります。

処理 計算量または追加メモリ
1回の find と併合 経路圧縮とサイズによる併合を合わせた償却計算量はO(α(N))
M本の辺の処理と最後のN回の根の確認 O((N + M) α(N))
parentsize、最後に作る集合 O(N)の追加メモリ

α(N)は、Nが増えても非常にゆっくり増える関数です。コードは辺をリストへ保存せず、入力を1本ずつ処理します。最後の集合も頂点数に比例する大きさなので、追加メモリ全体はO(N)に収まります。

削除を含む問題との境界

このコードが更新するのは、入力された辺を追加して集合を併合する状態です。すでに併合した集合から、任意の辺を削除して連結成分を分け直す処理は実装していません。

問題文に辺の削除が含まれる場合や、追加と削除が混ざる複数の遷移を扱う場合は、このコードをそのまま適用せず、操作列の条件に合う別手法を検討します。削除を扱う別手法の実装や、その問題の解説はこの範囲に含めません。

ABC284 Cのように、入力された辺を順に追加し、最後に連結成分数を集約する問題なら、親配列、経路圧縮、サイズによる併合という構成で答えまで追えます。