UnionFind

class UF:
    def __init__(self, size: int):
        self.size = size
        self.parent = [i for i in range(size)]
        self.component_sizes = [1 for _ in range(size)]
        self.num_components = size

    def find(self, p: int) -> int:
        """
        Find the root of the node and do path compression along the way.
        :param p: node which you want to find the root of
        :return: root node of the component to which `p` belongs
        """
        root = p
        while root != self.parent[root]:
            root = self.parent[root]
        current = p
        while current != root:
            nxt = self.parent[current]
            self.parent[current] = root
            current = nxt
        return root

    def connected(self, p: int, q: int) -> bool:
        return self.find(p) == self.find(q)

    def component_size(self, p: int) -> int:
        return self.component_sizes[self.find(p)]

    def unify(self, p: int, q: int):
        if self.connected(p, q):
            return
        r1 = self.find(p)
        r2 = self.find(q)

        smaller, larger = sorted([r1, r2], key=lambda x: self.component_sizes[x])
        self.component_sizes[larger] += self.component_sizes[smaller]
        self.parent[smaller] = larger
        self.component_sizes[smaller] = 0

        self.num_components -= 1