Union Find Data Structure using python also called DisJoint Sets

union_find.py

class UnionFind:
    def __init__(self):
        self.parent = dict()
        self.rank = dict()

    def find(self, item) -> int:
        origional_item = item

        parent = self.parent[item]
        while parent is not None:
            item = parent
            parent = self.parent[item]

        # Path compression
        if self.parent[origional_item] is not None and self.parent[origional_item] != item:
            self.parent[origional_item] = item

        return item

    def union(self, x, y) -> None:
        for element in [x, y]:
            if element not in self.parent:
                self.parent[element] = None
                self.rank[element] = 1

        parent_x = self.find(x)
        parent_y = self.find(y)

        if parent_x == parent_y:
            return

        rank_parent_x = self.rank[parent_x]
        rank_parent_y = self.rank[parent_y]

        if rank_parent_x > rank_parent_y:
            self.parent[parent_y] = parent_x
            self.rank[parent_x] = rank_parent_x + rank_parent_y
        else:
            self.parent[parent_x] = parent_y
            self.rank[parent_y] = rank_parent_x + rank_parent_y

    def get_no_of_islands(self):
        return sum(1 for _ in self.parent if self.parent[_] is None)


if __name__ == '__main__':

    DisJointSet = UnionFind()
    edges = [(1, 2), (2, 3), (4, 6)]
    for u, v in edges:
        DisJointSet.union(u, v)
    
    if DisJointSet.find(1) == DisJointSet.find(2):
        print("1 and 2 are in same set")
    
    print(f"No of islands in graph: {DisJointSet.get_no_of_islands()}")