#CP Union-Find木クラス Python

競プロメモ

ABC120 D問題にてUnion-Find木を扱う問題が出たのでUnion-Find木をクラスとして表現しておこうと思います。

パクり元

例のごとく、じゅっぴー氏のコードをパクります。いつもすいません。
juppy.hatenablog.com

ソースコード

class UnionFind:
    def __init__(self,n):
        self.par = [i for i in range(n)] #親
        self.rank = [1]*n #深さ
        self.size = [1]*n #size[i]:=iを根とするグループのサイズ
    
    def find(self,x):
        if self.par[x] == x:
            return x
        else:
            return self.find(self.par[x])
            
    def unite(self,x,y):
        x = self.find(x)
        y = self.find(y)
        if x != y:
            if self.rank[x] < self.rank[y]:
                self.par[x] = y
                self.size[y] += self.size[x]
            else:
                self.par[y] = x
                self.size[x] += self.size[y]
                if self.rank[x] == self.rank[y]:
                    self.rank[x] += 1
                    
    def same(self,x,y):
        return self.find(x) == self.find(y)
        
    def group_size(self,x):
        return self.size[self.find(x)]

使用法

※注意 ノード番号は0から始まります。

  • コンストラクタ: UnionFind(max_N)
    • max_N : ノードの個数(int)
  • xの親を見つける: find(x)
    • x : ノード番号(int)
  • xとyの属する集合を併合: unite(x,y)
    • x,y : いずれもノード番号(int)
  • xとyが同じ集合に属するかの判定: same(x,y)
    • x,y : いずれもノード番号(int)
    • 返り値: True if xとyが同じ集合に属する else False
  • xが属する集合の大きさ: group_size(x)
    • x : ノード番号(int)
    • 返り値: xが属する集合の大きさ(int)

使用例

### 定義は省略 ###
uf = UnionFind(7)
#A(0,1,4) B(2) C(3,5,6)のグループを作成
uf.unite(0,1)
uf.unite(0,4)
uf.unite(3,5)
uf.unite(5,6)
#1と2,3と5は同じ集合かの判定
print(uf.same(1,2))
print(uf.same(3,5))
#AとBの併合
uf.unite(4,2)
#1と2は同じ集合かの判定
print(uf.same(1,2))
#1の属する集合の大きさ
print(uf.group_size(1))

結果

False
True
True
4