蟻本初級編攻略 - 2-4 Union Find木 前編

「ある集合の二つの要素が繋がっているか」を効率的に判定できるデータ構造であるUnion Find木の話。

非常に簡単でエレガントな実装のわりに非常に強力な概念だという印象で、個人的にはとても好き。

蟻本の例題が一番特殊で面白いので後回し。まずはAtCoder Typical Contestから：

AtCoder Typical Contest 001 B問題 Union Find

atc001.contest.atcoder.jp

atc001.contest.atcoder.jp

ザ・Union Find典型。Union Findが逐次更新できるデータ構造で、途中のいつでも「現在明らかになった情報を元に二つの要素が繋がっているか判定できる」という点も強調されている。

Python実装はこんな感じ：

n, q = map(int, input().split())
xs = [tuple(map(int, input().split())) for _ in range(q)]
 
union = {x:x for x in range(n)}
def root(n):
    if union[n] == n:
        return n
    union[n] = root(union[n])
    return union[n]
 
for p, a, b in xs:
    if p == 0:
        union[root(a)] = root(b)
    else:
        print("Yes" if root(a) == root(b) else "No")

やはり非常に簡潔に済むのがうれしい。

Union Findはデータをdictionaryに持たせ、そのdictionaryをもとにある要素が繋がっているグループのルートノードを返す関数と、二つのグループをつなげるイディオムを提供する。二つの要素が繋がっているかを判定するのは、その二つの要素の属するグループのルートノードが一致するかを調べるだけ。

Union Findは気をつけないとroot関数が最悪O(n)かかってしまう。それを回避するメジャーな最適化が二つある。ここらへんはAtCoderのUnion FindについてのSlideShareに詳しい：

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

www.slideshare.net

両方の最適化を使うと計算量のオーダーが逆アッカーマン関数（ほぼ定数・・・）になるということだが、一つだけでもO(logN)になるうえ、実装が簡単で定数倍ははやい。

というわけで私は「ルートノードを探る度に枝をすべてルートノード直下に貼り直す」という最適化のみ実装している。

つまりroot関数のナイーブな実装：

def root(n):
    if union[n] == n:
        return n
    return root(union[n])

それをこう変えている：

def root(n):
    if union[n] == n:
        return n
    union[n] = root(union[n]) #ここ
    return union[n]

これだけでO(N)がO(log N)になるのはうれしい。

そういえば昨日PEP572が受理されたので将来こう書けるようになる：

def root(n):
    if union[n] == n:
        return n
    return union[n] := root(union[n]) #ここ

まあこれは大した違いじゃないけど・・・。

競プロ時にPythonでClassを使わずにデータ構造を実装することについて

ちなみに私は個人的に競技プログラミングでPythonを書く場合、Union Findくらいだったらインターフェースつけてclassにwrapするより、もとのデータ構造をむき出しにしておきたい。

ということをちょっと乱暴に呟いた：

競プロで絶対class書きたくないマン

— zehnpaard (@zehnpaard) 2018年7月5日

そしたら速攻でツッコミが入った：

タプル使うよりClass使ったほうがよくないですか？

— すとまと (@stmtk_01) 2018年7月5日

（次回出てくる重みつきUnion Find木の話）

もうちょっと丁寧に自分の考えを説明するとこんな感じ：

競プロだと「0から読み直して頭の中でフローを再構築できるコード」である必要はなくて「ある程度頭の中でコンテキストが出来上がっている中でフローが追いやすくエラーが見つけやすいコード」が最適なので、記述の仕方もやっぱりけっこう違うなーと思う

— zehnpaard (@zehnpaard) 2018年7月5日

（最初のnamedtupleについての言及はプロダクションコードだったら、の話）

さらにいうとデータ構造がむき出しのほうが短期的にはいじりやすい。微妙にデータ構造に持たせる情報や算出させる情報を変えたいときに、すくなくともABCレベルの競技プログラミングだと変更箇所はかなり限られているから使われているその場で変えてしまったほうが楽。

数百行以上のコードのいろんなところで使われているなら状況は逆転する。やはりプロダクションコードと競技プログラミングでは細かい作法がけっこう違ってくるのは自然だと思う。

次回予告

つらつらと書いていたら長くなったので記事を二回にわける。

次回はAtCoder ABC/ARCからUnion Findを使ったD問題４つと蟻本の例題の解法を説明する。