Pythonで理解する蟻本「2-4 Union-Find木」(p.81)

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 Union-Find木」(p.81)
のコードをPythonで書き直したものとなっています。

Union-Find木の実装

コード

MAX_N = 10 ** 8

par = [0] * MAX_N     # 親
rank = [0] * MAX_N    # 木の深さ

# n要素で初期化
def init(n):
    for i in range(n):
        par[i] = i
        rank[i] = 0

# 木の根を求める
def find(x, par):
    if par[x] == x:
        return x
    else:
        par[x] = find(par[x], par)
        return par[x]

# xとyの属する集合を併合
def unite(x, y, par, rank):
    x = find(x, par)
    y = find(y, par)
    if x == y:
        return
    if rank[x] < rank[y]:
        par[x] = y
    else:
        par[y] = x
        if rank[x] == rank[y]:
            rank[x] += 1

# xとyが同じ集合に属するか否か
def same(x, y, par, rank):
    return find(x, par) == find(y, par)

実行例

コード

init(100)                       # 要素数が100のUnion-Find木を作る
print(same(0, 1, par, rank))    # 0と1が同じ集合に属するかを確認する
unite(0, 1, par, rank)          # 0と1の属する集合を併合する
print(same(0, 1, par, rank))    # 0と1が同じ集合に属するかを確認する

実行結果

False
True

classを使ったUFTの実装

上記の実行例を見ると、関数による実装では少し扱いづらいことが分かります。

そのため、ここではclassによる実装も載せておこうと思います。

（実際にこのコードを使ってUnion-Find木の問題を解けることは確認していますが、もしバグを発見された場合は僕のTwitterアカウント（クルトン (@kuruton456) | Twitter）に連絡して頂けると助かります。）

コード

class UnionFindTree():
    # n要素で初期化
    def __init__(self, n):
        self.n = n
        self.par = list(range(n))
        self.rank = [0] * n

    # 木の根を求める
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # xとyの属する集合を併合
    def unite(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    # xとyが同じ集合に属するか否か
    def same(self, x, y):
        return self.find(x) == self.find(y)

実行例

コード

UFT = UnionFindTree(100)    # 要素数が100のUnion-Find木を作る
print(UFT.same(0, 1))       # 0と1が同じ集合に属するかを確認する
UFT.unite(0, 1)             # 0と1の属する集合を併合する
print(UFT.same(0, 1))       # 0と1が同じ集合に属するかを確認する

実行結果

False
True