Union Findをpythonで実装して関数化する

hirohirohirohiros.hatenablog.com

概要

　Union Findはグループ分けを管理するデータ構造です．要素aと要素bが同じグループに属しているかそうでないか判別することが出来ます．
　また，複数のグループを一つのグループにくっつけることも出来ます．このようにUnion Findは複数機能を有しており，同じリストを使って管理するので，クラスを使って実装します．コードは蟻本を，クラスを使うというアイデアはnkmkさんのサイトhttps://note.nkmk.me/python-union-find/を参考にしました．

アルゴリズムの概要

併合，判定

　木構造を使います．n個の要素に対してUnion Findを使うとき，始めn個のノードを作ります．辺はありません．
　グループをくっつける(unite)には片方のグループの親からもう片方のグループの親に辺を張ればグループが一つになります．
　ノードaとノードbが同じグループに属してるか判定するにはそれぞれのグループの親を調べます．同じグループなら同じ親になるはずです．
　このことから，各ノードには親を記憶しておく必要があると分かります．親の管理にはリストを使います．リストのi番目はノードiの親の番号を記録すれば良いです．そして，リストのi番目の値とj番目の値が同じなら，ノードiとノードjは同じグループに属していることになります．uniteするにはどちらかの親をどちらかの親の値に書き換えれば良いです．どちらのノードの親を書き換えれば良いのでしょうか？より効率よくするためにrankという概念を導入します．

rank

　rankはグループの木の高さを表します．左のグループはrank1で右のグループはrank2です．rankが低い方が親を探すときに掛かる回数が少ないことが分かります．
この二つのグループをuniteする時どちらのグループをどちらのグループにuniteすれば良いでしょうか？rankの小さい方をrankの大きい方にくっつければrankが増えずに済む事が分かります．
よってグループにrankを記録することでより効率よくuniteが出来る事が分かります．

pythonコード

class Unionfind:
    """ Unionfind

    ノードのグループ化, 同一グループに属しているかの判定
    計算量: O(α(n)) log n より早い

    """
    def __init__(self, n):
        """

        Args:
            n (int): ノードの総数
        """
        self.par = [i for i in range(n)]
        self.rank = [0]*n

    def find(self, x):
        """親が何か探す

        Args:
            x (int): 親を探したいノード番号

        Returns:
            int: 親のノード番号
        """
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    def unite(self, x, y):
        """グループの併合

        Args:
            x, y (int): 併合したいノードの番号
        """
        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
    def same(self, x, y):
        """同じグループに属しているか判定

        Args:
            x, y (int): 同じグループに属しているか判定したいノードの番号

        Returns:
            bool: 同じグループならTrue, そうでないならFalse 
        """
        return self.find(x) == self.find(y)