Union-Find with Path Compression and Rank

from __future__ import annotations


class DisjointSet:
    def __init__(self, items: list[str]) -> None:
        self.parent = {item: item for item in items}
        self.rank = {item: 0 for item in items}
        self.components = len(self.parent)

    def find(self, item: str) -> str:
        root = item
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[item] != item:
            parent = self.parent[item]
            self.parent[item] = root
            item = parent
        return root

    def union(self, left: str, right: str) -> bool:
        root_left = self.find(left)
        root_right = self.find(right)
        if root_left == root_right:
            return False
        if self.rank[root_left] < self.rank[root_right]:
            root_left, root_right = root_right, root_left
        self.parent[root_right] = root_left
        if self.rank[root_left] == self.rank[root_right]:
            self.rank[root_left] += 1
        self.components -= 1
        return True

    def connected(self, left: str, right: str) -> bool:
        return self.find(left) == self.find(right)


def main() -> None:
    network = DisjointSet(["a", "b", "c", "d", "e"])
    for left, right in [("a", "b"), ("b", "c"), ("d", "e")]:
        print(left, right, network.union(left, right))
    print(network.connected("a", "c"), network.connected("a", "d"))
    assert network.components == 2
    assert not network.union("a", "c")


if __name__ == "__main__":
    main()