Merging friend groups

Union-Find (Disjoint Set Union)

The analogy

Everyone starts in their own group. When two people become friends you merge their whole groups, and to check whether two people are already connected you just ask "do you have the same group leader?" Each person remembers who to ask, and after a while everyone points almost directly at the leader.

Visualizer

Merging the forest

step 1 / 8
ABCDE

Five people, five separate groups. Everyone is their own leader.components: 5

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]  # compress
            x = self.parent[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False                  # already together -> cycle
        if self.size[ra] < self.size[rb]:  # union by size
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True
Check yourself

union(a, b) returns False. What does that tell you?