【Python】例題でUnion-Find木を理解する

f:id:mizukawa815:20211014085742p:plain
この記事では、PythonによるUnion-Find木の実装と解説をしています。

Union-Find木とはなんぞやと思っている方もこの記事でその仕組みを理解していきましょう。

1. はじめに
- 1.1. Union-Find木とは
2. 例題
4. 参考・参照

1. はじめに

1.1. Union-Find木とは

素集合データ構造（そしゅうごうデータこうぞう、英: disjoint-set data structure）は、データの集合を素集合（互いにオーバーラップしない集合）に分割して保持するデータ構造。このデータ構造に対する以下の2つの便利な操作をUnion-Findアルゴリズムと呼ぶ。

　・Find: 特定の要素がどの集合に属しているかを求める。2つの要素が同じ集合に属しているかの判定にも使われる。
　・Union: 2つの集合を1つに統合する。

これら2つの操作をサポートしているため、素集合データ構造は「Union-Findデータ構造」あるいは「Merge-Find集合」とも呼ばれる。これら以外の重要な操作として MakeSetがある。これは、与えられた1つの要素だけを格納した集合（シングルトン）を作成する。これら3つの操作により、様々な実用的分割問題を解くことができる（「応用」の節を参照）。

引用元：素集合データ構造 - Wikipedia

Union-Find木は、競技プログラミングなどでも度々登場するアルゴリズムです。

一見難しそうに感じますが、やっていることは単純なのでこの機会に習得しておくことをおすすめします。

文字だけではなんのこっちゃですが、以下のスライドではイラストを交えて詳しく解説してくれています。

Union find(素集合データ構造) from AtCoder Inc.

では、さっそく例題をPythonで実装していきます。

2. 例題

2.1. 【AtCoder Typical Contest 001】B - Union Find

N 頂点の、単純とは限らない無向グラフを考えます。初期状態では、頂点のみが存在し、辺は全く存在せず、全ての頂点が孤立しているとします。以下の 2 種類のクエリが、Q 回与えられます。
　・連結クエリ：頂点 A と、頂点 B を繋ぐ辺を追加します。
　・判定クエリ：頂点 A と、頂点 B が、連結であるかどうか判定します。連結であれば Yes、そうでなければ No を出力します。
連結であるとは、頂点 A から頂点 B まで辺をたどって到達可能であることを意味します。 A と B が同じ頂点の場合、連結であると定義します。グラフは無向であるため、連結クエリによって頂点 A,B 間の辺が追加されると、A から B へも B から A へも辿れるようになります。

　・P=0 の時、そのクエリが連結クエリであることを表す。
　・P=1 の時、そのクエリが判定クエリであることを表す。

2.2. 実装

n, q = map(int, input().split()) # N, Qの入力
pair = [i for i in range(n)] # 親番号の初期化

# 根を見つける
def find(x):
    if(pair[x] == x):
        return x
    else:
        pair[x] = find(pair[x])
        return pair[x]
        
# x, yの親の判定
def check(x, y):
    return find(x) == find(y)

# x, yの結合
def union(x, y):
    if find(x) == find(y):
        return
    
    pair[x] = y

# 実行
def main():
    for i in range(q):
        p, a, b = map(int, input().split()) # クエリQと操作する要素A, Bの入力
        if p == 0:
            union(a, b)
        else:
            if check(a, b):
                print('Yes')
            else:
                print('No')

if __name__ == "__main__":
    main()

2.2. 解説

それぞれの関数の役割を順番に解説していきます。

n, q = map(int, input().split()) # N, Qの入力
pair = [i for i in range(n)] # 親番号の初期化

まず初めに頂点の数Nとクエリを与えられる回数Qを入力します。
それぞれの要素の親の番号をpairに格納、初期状態ではそれぞれが根です。

# 根を見つける
def find(x):
    if(pair[x] == x):
        return x
    else:
        pair[x] = find(pair[x])
        return pair[x]

find関数では、それぞれの要素の根を見つけます。
例えば、『子→親』という表記で『⑤→②』『②→②』という関係があれば、「5の親は2、2の親は2」であるため、2は根となります。

# x, yの親の判定
def check(x, y):
    return find(x) == find(y)

判定クエリでそれぞれの要素が同じ集合に属するかを判定するために使用する関数です。
戻り値としてBoolean型を返します。
例えば、「x, yがそれぞれ同じ親を持っていれば同じ集合の要素である」と言え、戻り値としてTrueを返します。

# x, yの結合
def union(x, y):
    if find(x) == find(y): # 同じ集合であれば何もしない
        return
    
    pair[x] = y

結合クエリでそれぞれの要素を辺で繋ぐために使用する関数です。
例えば、「x=2とy=5を結合する場合、2の子要素に5を設定」します。
x, yの親要素が同じ場合は既に結合済みのため何もしません。

# 実行
def main():
    for i in range(q):
        p, a, b = map(int, input().split()) # クエリQと操作する要素A, Bの入力
        if p == 0:
            union(a, b)
        else:
            if check(a, b):
                print('Yes')
            else:
                print('No')

if __name__ == "__main__":
    main()