union-find.py

class UnionJoin:
    def __init__(self,n):
        # parent node point to itself
        self.parent = [i for i in range(n)]
        self.size = [0 for _ in range(n)]
        self.count = n

    def union(self,p, q):
        # parent of p is parent of parent of q
        parent_p = self.find_parent(p)
        parent_q = self.find_parent(q)
        if parent_q == parent_p:
            return

        if self.size[p] > self.size[q]:
            self.size[q] += 1
            self.parent[q] = p
        else:
            self.size[p] += 1
            self.parent[p] = q
        self.count -= 1



    def isConnected(self,p, q):
        parent_p = self.find_parent(p)
        parent_q = self.find_parent(q)

        return parent_q == parent_p

    def find_parent(self, p):
        while self.parent[p] != p:

            # Path compression
            self.parent[p] = self.parent[self.parent[p]];
            p = self.parent[p]
        return p


uf = UnionJoin(5)

uf.union(0, 1)
uf.union(1, 2)
uf.union(1, 3)
uf.union(1, 4)

print(uf.size)
print(uf.parent)
assert uf.isConnected(0, 2) == True
assert uf.isConnected(1, 4) == True
print(uf.parent)

assert uf.count == 1