Python Implementation of Union Find Tree

union_tree.py

# =============
# UnionFind tree
# =============
class UnionFindTree():
    def __init__(self, N):
        self.parent = [i for i in range(N)]
        self.rank = [0 for _ in range(N)]

    def root(self, x):
        if x == self.parent[x]:
            return x
        self.parent[x] = self.root(self.parent[x])
        return self.parent[x]

    def same(self, x, y):
        return self.root(x) == self.root(y)

    def unite(self, x, y):
        x = self.root(x)
        y = self.root(y)

        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.parent[x] = y
        else:
            self.parent[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1


if __name__ == '__main__':
    # test UnionFind Tree
    uft = UnionFindTree(N=5)
    uft.unite(0, 1)
    uft.unite(2, 3)
    uft.unite(0, 4)
    assert uft.same(0, 2) == False
    assert uft.same(0, 1) == True
    assert uft.same(1, 4) == True
    assert uft.root(1) == 0
    assert uft.root(0) == 0

reference:
https://www.slideshare.net/chokudai/union-find-49066733