class UnionFind: # Constructor def __init__(self, n): self.parent = list(range(n)) self.rank = [1]*n # Function to find which subset a particular element belongs. def find(self, i): if self.parent[i] != i: self.parent[i] = self.find(self.parent[i]) return self.parent[i] # Function to join two subsets into a single subset. def union(self, x, y): root_x, root_y = map(self.find, (x, y)) if root_x == root_y: return small, big = sorted([root_x, root_y], key=lambda z: self.rank[z]) self.parent[small] = big self.rank[big] += self.rank[small]