Union FindをPythonで実装して基本的な事を理解する

Union Findとは？

ある配列のグループ分けを効率的に管理するデータ構造となります。

例えば、出席番号1~6が付いた6人がおり、友達の友達は友達であるという条件のもと、(1, 2)、(3, 4)、(5, 6)、(5, 2)は友達という条件が与えられた時に、1番の人は何人友達がいるか？などを計算したい場合に使えます。

Union Findでは下記のように、人をノード、友達関係をエッジとして、集合を木構造で管理することで効率的なデータ管理を行います。

色々と拡張機能はあるようですが、私自身も勉強中のため、今回は基本的な内容を自分用にまとめます。

実装例

一例として下記のように実装できます。

class UnionFind:
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return
        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

class UnionFind:

def __init__(self, n):

self.n = n

self.parents = [-1] * n

def find(self, x):

if self.parents[x] < 0:

return x

else:

self.parents[x] = self.find(self.parents[x])

return self.parents[x]

def union(self, x, y):

x = self.find(x)

y = self.find(y)

if x == y:

return

if self.parents[x] > self.parents[y]:

x, y = y, x

self.parents[x] += self.parents[y]

self.parents[y] = x

def size(self, x):

return -self.parents[self.find(x)]

関数名が表す通り、unionメソッドで2つの要素を同じ集合にグループ化して、sizeメソッドで各ノードが属する集合の要素数を確認することができます。

簡単な例で理解

例えば、出席番号1から6の6人の中で、友達関係を表す条件が4つが順番に与えられた時に、1番の人が属してる集合の要素数を取得するケースを考えます。

友達関係のグループ化と、サイズを出力する処理は上で定義したクラスを使用することで、下記のように書けます。

# Nはノード数, Mは条件数(友達関係)
N = 6
M = 4

UF = Union_Find(N) # ノード数で初期化
for _ in range(M):
    x, y = inputmap() # 友達関係を取得(x=1,y=2)
    UF.union(x-1, y-1) # 同じ集合にする (0-index)

print(UF.size(0)) # 1番目が属する集合の要素数を取得

# Nはノード数, Mは条件数(友達関係)

N = 6

M = 4

UF = Union_Find(N) # ノード数で初期化

for _ in range(M):

x, y = inputmap() # 友達関係を取得(x=1,y=2)

UF.union(x-1, y-1) # 同じ集合にする (0-index)

print(UF.size(0)) # 1番目が属する集合の要素数を取得

友達関係が順番に(1, 2)、(3, 4)、(5, 6)、(2, 5)と与えられたとして、処理がどのように行われているかを確認します。

最初に、ノード数と同じ長さを持つリスト(self.parents)を用意して-1で初期化を行います。このリストは親ノードor根ノードの場合ノード数の情報を持たせます。

根ノードとは、木構造を作成した時に親ノードを持たないノード、親ノードとは自身が結合しているノードの中で根ノードに1つ近いノードの事を指します。初期化した時点では、全てが根ノードとなっています。

self.parentsには、根ノードに-(ノード数)が入り、それ以外には、親ノードのindexが入ります。例えば、1つ目の条件である(1,2)という条件から、0-indexに注意して、UF.union(0,1)という処理を実施すると下記のように結合され、self.parentsの値が更新されます。

同様に(3, 4), (5, 6)という条件を使用すると下記のように処理されます。

3つの集合が作成できて、根ノードには、-(ノード数)が格納され、それ以外のノードには、親ノードのindexの情報を格納することができました。

木同士の結合と経路圧縮

続いて、条件(2,5)のような木同士の結合を考えます。これまでと同じようにUF.union(1, 4)でグループ化を行うと下記のように処理されます。

条件(2, 5)を与えていますが、ノード5はノード2を参照するのではなく、ノード1を参照するようにリストが更新されます。

これは、結合先の根ノードを探索する際、走査したノード全てを、根ノードのindexを参照するように書き換える短縮経路という工夫によるものとなります。

例えば、上のような1~4ノードがある場合、4の親は3→3は根ノードではないので親を探すと2→2は根ノードではないので親ノードを探すと1→2,3,4の親ノードを全て1にするように書き換えるという流れになります。

実装上は、下記のように、各ノード毎に根ノードを再帰的に探索して、最終的に取得した根ノードのindexを、走査したノードのself.parentsに格納しています。また、ノード6は結合した後に走査していないため、この段階では根ノード1を参照しておらず、ノード5に紐づく形で管理されます。

こうすることで、例えば木が長くても、根ノードを探索する場合、一度走査したノードに対してはO(1)で取得することができます。

また、そのほかの工夫として、マージする際に大きな木に小さな木を結合するように実装が行われています。

これについては、自分の理解では、走査する回数を減らすため(木が短い方の根ノードを書き換える方が処理量が少ないから)であると思っています(少し調べたましたが、理由についてはわからなかった…)

Union Find を使う問題

精進中に解いた問題を随時追加していきます。

AtCoder Beginner Contest 177 D – Friends

AtCoder Beginner Contest 206（Sponsored by Panasonic）D – KAIBUNsyo