[競プロ用]UnionFind木まとめ

いつ使うのか

• グラフが同一の根を持つか → 全てのor特定の点同士が繋がっているかを考える問題
• 複数の要素が同一の集合に含まれているかを考える問題

何がいいのか

O(α(N))の非常に少ないオーダー。
α(n)はアッカーマン関数A(n, n)の逆数。

テンプレート

とても便利なライブラリを作ってくださっている。。
https://note.nkmk.me/python-union-find/
https://github.com/nkmk/python-snippets/blob/ec0e770a72135b12840ea02f1e74101785faf700/notebook/union_find.py#L1-L46

union()で要素の併合→same()やgroup_count()で求めたい結果を得る。
難しくなさそう

使用例

ARC032 B - 道路工事

def main():
    N, M = map(int, input().split())
    uf = UnionFind(N)
    for _ in range(M):
        aa, bb = map(lambda x: int(x) - 1, input().split())
        uf.union(aa, bb)
    print(uf.group_count() - 1)

ABC075 C - Bridge

def main():
    N, M = map(int, input().split())
    a, b = [], []
    for _ in range(M):
        aa, bb = map(int, input().split())
        a.append(aa - 1)
        b.append(bb - 1)
    count = 0
    for i in range(M):
        uf = UnionFind(N)
        for j in range(M):
            if j == i:
                continue
            uf.union(a[j], b[j])
        if uf.group_count() != 1:
            count += 1
    print(count)

ABC157 D - Friend Suggestions

def main():
    N, M, K = map(int, input().split())
    connection = [[] for _ in range(N)]
    uf = UnionFind(N)
    for i in range(M):
        a, b = map(lambda x: int(x) - 1, input().split())
        connection[a].append(b)
        connection[b].append(a)
        uf.union(a, b)

    for i in range(K):
        c, d = map(lambda x: int(x) - 1, input().split())
        if uf.same(c, d):
            connection[c].append(d)
            connection[d].append(c)

    ans = []
    for i in range(N):
        ans.append(uf.size(i) - len(connection[i]) - 1)
    L = [str(i) for i in ans]
    print(" ".join(L))

ARC097 D - Equals

こういう問題でもUnionFind。
i = p[i]となるかどうかがuf.same(i, p[i])と置き換えられる。

def main():
    N, M = map(int, input().split())
    p = [int(i) - 1 for i in input().split()]
    uf = UnionFind(N)
    for _ in range(M):
        x, y = map(lambda i: int(i) - 1, input().split())
        uf.union(x, y)

    count = 0
    for i in range(N):
        if uf.same(i , p[i]):
            count += 1
    print(count)

参考

• https://www.slideshare.net/chokudai/union-find-49066733
• https://note.nkmk.me/python-union-find/
• https://github.com/nkmk/python-snippets/blob/ec0e770a72135b12840ea02f1e74101785faf700/notebook/union_find.py#L1-L46

Appendix

map関数の第一引数にはlambda式を使える。
これまで入力を文字列で受けてintに変えるくらいにしか使っていなかったが、リストの0-indexedか1-indexedかの問題で入力値をマイナス1する時にlambda式にすると一括してできそう。

map(lambda i: int(i) - 1, input().split())