Union Find

Advent Of Code をやっていて、エッジのリストから、何個のグラフができるかを求めたくなりました。

pairs = [(0, 1), (1, 2), (3, 6), (4, 5), (5, 6), (7, 8)]

# pairs は何個のグラフに分けられるか？
# 上記の例だと以下の3つ
# - 0, 1, 2
# - 3, 4, 5, 6
# - 7, 8

とりあえず素朴にやるなら以下だと思います。

• (0, 1) を取り出す。
  • まず、0 がないので登録。探しやすいように自分自身も入れる。
  • 1 を 0 の辞書に追加
• (1, 2) を取り出す。
  • 辺のどちらかがすでに値として保存されていないかを調べる
    • すでに辞書に登録されている値の中に、1がないかを探す。 -> ある
    • すでに辞書に登録されている値の中に、2がないかを探す。 -> ない
  • 一方だけが、「ある」なので、その値に追加する。すでに登録されているものは除く。

仮に、取り出した (a, b) の両方が、「ある」で、異なるキーの中にある場合は、2つのキーをくっつける。

d = {
    0: [0, 1, 2]
    3: [3, 6] # (5, 6) が、こことその下に一致する
    4: [4, 5]
}

d = {
    0: [0, 1, 2]
    3: [3, 6, 4, 5] # くっつけていらなくなったものを消す
}

これはこれでいいかもしれませんが、もう少しいい方法あるのでは？と思って調べたところ、Union Find というのがちょうど、 「データの集合を素集合に分割して保持」・「特定の要素が、どの集合にあるかを特定」に使えるということを知りました。

def union_find(edges):
    nodes = set()
    for a, b in edges:
        nodes.add(a)
        nodes.add(b)
    
    parent = {n: n for n in nodes} # {0: 0, 1: 1, 2: 2, ...} を作る
    
    # ルートを見つける
    def find(x):
        if parent[x] != x:
            parent[x] = find(parent[x]) # もし推移的 ({1: 2, 2: 3, 3: 3}) になっていたら、直接ルートと繋げる ({1: 3, 2: 3, 3: 3})
        return parent[x]
    
    # 集合の統合
    def union(x, y):
        # x を要素として持つ集合と、 y を要素として持つ集合を統合する。
        # x のルート の 親 を y のルート に変更する。
        parent[find(x)] = find(y)
    
    # はじめはノード1個だけの素集合がノードの個数分
    # 集合の統合をエッジの数だけ繰り返す。
    for a, b in edges:
        union(a, b)

find の箇所（経路圧縮）がとても頭良すぎるな…と感動を覚えました。 こういうの知るために競プロ周りの問題を改めてやってみても良いなと思いました。

以上！

おまけ

これはこの記事を書いた後にAIに作ってもらったビジュアライザーです。

私にとってのAIのワクワクする使い方ってこれなんですよね。動く絵本…人の書いたものをインタラクティブな読み物に変えてほしい という思いがあります。