PythonでUnion-Findアルゴリズムを実装する

けんちょん本（「アルゴリズムとデータ構造」）の第11章で扱われているUnion-Findについて、Pythonで実装してみました。

前提知識：素集合データ構造に対する操作

素集合データ構造とは、互いに重なり合うことが無い集合を持つデータ構造（おおざっぱにいうときれいに分類されているということ）。このデータ構造に対して、Union（2つの集合を1つに統合する）とFind（特定の要素がどの集合に属しているかを調べる）という操作を行うアルゴリズムをUnion-Findアルゴリズムと呼びます。この2操作に、MakeSet（集合を作成する）を加えることにより、さまざまな分割問題をとけるようなツールとして利用できます。

Pythonで実装してみた

けんちょん本とはちょっと実装方針を変えています。

• 素集合データ構造（DisjointSet）の中にUnion-Findを含めた。
• Root→Find、Unite→Unionにメソッド名を変えた。

class DisjointSet:
    def __init__(self) -> None:
        self.parent = {}
        self.size = {}

    def makeSet(self, set):
        for e in set:
            self.parent[e] = e
            self.size[e] = 1

    def Find(self, x):
        if self.parent[x] == x:
            return x
        else:
            # 経路圧縮
            self.parent[x] = self.Find(self.parent[x])
            return self.parent[x]

    def Union(self, x, y):
        a = self.Find(x)
        b = self.Find(y)

        if a == b:
            return False

        # union by sizeで効率化
        if self.size[a] < self.size[b]:
            a, b = b, a

        self.parent[b] = a
        self.size[a] += self.size[b]
        return True

    def issame(self, x, y):
        return self.Find(x) == self.Find(y)

    def getsize(self, x):
        return self.size[self.Find(x)]

def printSet(set, ds):
    return [ds.Find(i) for i in set]

def test_unionfind():
    ds = DisjointSet()
    ds.makeSet([1, 2, 3, 4, 5, 6])
    ds.Union(1, 2)
    ds.Union(2, 3)
    ds.Union(5, 6)

    assert ds.issame(1, 3) == True
    assert ds.issame(2, 5) == False

    ds.Union(1, 6)

    assert ds.issame(2, 5) == True

参考

• Disjoint–Set Data Structure (Union–Find Algorithm)