# Title: Union-find
# Heading: Python Program for Union-Find (Disjoint Set)
# Level: advanced
# Tags: union-find, disjoint set, dsu, path compression, connected components
# Time: O(α(n))
# Space: O(n)
# About: Puts friends into groups with union-find, where merging two groups and asking if two people share one each take almost constant time.

class UnionFind:
    """Items in groups that can be merged and looked up."""

    def __init__(self, items):
        self.parent = {item: item for item in items}  # all alone
        self.size = {item: 1 for item in items}

    def find(self, item):
        """The root of item's group."""
        if self.parent[item] != item:
            # Path compression: point item straight at the root
            self.parent[item] = self.find(self.parent[item])
        return self.parent[item]

    def union(self, a, b):
        """Merges the groups of a and b, smaller under larger."""
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return
        if self.size[root_a] < self.size[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]


friends = UnionFind(["Ann", "Ben", "Cal", "Dee", "Eve", "Fay"])
friends.union("Ann", "Ben")
friends.union("Cal", "Dee")
friends.union("Ben", "Dee")
friends.union("Eve", "Fay")
print("Ann and Dee:", friends.find("Ann") == friends.find("Dee"))
print("Ann and Eve:", friends.find("Ann") == friends.find("Eve"))