Union-Find木とは

Union-FindまたはDisjoint Setとは、蟻本の言葉を借りると「グループ管理のためのデータ構造」。

・要素aと要素bが同じグループかどうか

・aが属するグループと、bが属するグループを併合(unite)する

ということをしたいときに使います。

Union-Find木（深さバージョン）の実装

蟻本に書いてあるとおりの実装。

グループの併合のとき、木の深さが深いものに浅いものを併合する。

あまり使わない。競プロでは後述のサイズバージョンばかり使用している。

class UnionFindVerDepth():
    def __init__(self, N):
        """ N個のノードのUnion-Find木を作成する """
        # 親の番号を格納する。自分が親だった場合は、自分の番号になり、それがそのグループの番号になる
        self._parent = [n for n in range(0, N)]
        # グループの深さ
        self._depth = [1] * N

    def find_root(self, x):
        """ xの木の根(xがどのグループか)を求める """
        if self._parent[x] == x: return x
        self._parent[x] = self.find_root(self._parent[x]) # 縮約
        return self._parent[x]

    def unite(self, x, y):
        """ xとyの属するグループを併合する """
        gx = self.find_root(x)
        gy = self.find_root(y)
        if gx == gy: return

        # 小さい方を大きい方に併合させる（木の偏りが減るので）
        if self._depth[gx] < self._depth[gy]:
            self._parent[gx] = gy
        else:
            self._parent[gy] = gx
            if self._depth[gx] == self._depth[gy]: self._depth[gx] += 1

    def get_depth(self, x):
        """ xが属するグループの深さを返す """
        return self._depth[self.find_root(x)]

    def is_same_group(self, x, y):
        """ xとyが同じ集合に属するか否か """
        return self.find_root(x) == self.find_root(y)

    def calc_group_num(self):
        """ グループ数を計算して返す """
        N = len(self._parent)
        ans = 0
        for i in range(N):
            if self.find_root(i) == i:
                ans += 1
        return ans

Union-Find木（サイズバージョン）の実装

Union-Findのサイズバージョン。

グループの併合のとき、グループのサイズ（ノードの個数）が大きいものに、小さいものを併合させる。

競プロでは大抵、ノードの個数やグループ数の最小値とかを求めるので、こちらを使うことが多い。

class UnionFindVerSize():
    def __init__(self, N):
        """ N個のノードのUnion-Find木を作成する """
        # 親の番号を格納する。自分が親だった場合は、自分の番号になり、それがそのグループの番号になる
        self._parent = [n for n in range(0, N)]
        # グループのサイズ（個数）
        self._size = [1] * N

    def find_root(self, x):
        """ xの木の根(xがどのグループか)を求める """
        if self._parent[x] == x: return x
        self._parent[x] = self.find_root(self._parent[x]) # 縮約
        return self._parent[x]

    def unite(self, x, y):
        """ xとyの属するグループを併合する """
        gx = self.find_root(x)
        gy = self.find_root(y)
        if gx == gy: return

        # 小さい方を大きい方に併合させる（木の偏りが減るので）
        if self._size[gx] < self._size[gy]:
            self._parent[gx] = gy
            self._size[gy] += self._size[gx]
        else:
            self._parent[gy] = gx
            self._size[gx] += self._size[gy]

    def get_size(self, x):
        """ xが属するグループのサイズ（個数）を返す """
        return self._size[self.find_root(x)]

    def is_same_group(self, x, y):
        """ xとyが同じ集合に属するか否か """
        return self.find_root(x) == self.find_root(y)

    def calc_group_num(self):
        """ グループ数を計算して返す """
        N = len(self._parent)
        ans = 0
        for i in range(N):
            if self.find_root(i) == i:
                ans += 1
        return ans

使い方

N = 10
uf = UnionFindVerSize(N)

uf.unite(0, 1)
uf.unite(1, 2)
uf.unite(2, 3)
uf.unite(4, 5)
uf.unite(4, 6)

# {0, 1, 2, 3}, {4, 5, 6}, {7}, {8}, {9} の5グループ

a = uf.find_root(2)
print(a) # 0

a = uf.get_size(2)
print(a) # 4

a = uf.is_same_group(3, 4)
print(a) # False

a = uf.calc_group_num()
print(a) # 5