Union-find

Union-find, also called a disjoint set, keeps track of which items are in the same group. It does two things: union joins two groups, and find tells you which group something's in. Both are close to O(1). It shines when connections keep arriving and you need to know who's connected after each one.

Every item points to a parent, and the item at the top of each chain represents the group. To join two groups, point one group's top at the other's. Two tricks keep the chains short: attach the smaller group under the bigger one, and while finding, point items closer to the top as you go.

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

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

    def union(self, a, b):
        a, b = self.find(a), self.find(b)
        if a == b:
            return False
        if self.size[a] < self.size[b]:
            a, b = b, a
        self.parent[b] = a  # Smaller group joins the bigger one.
        self.size[a] += self.size[b]
        self.groups -= 1
        return True

union returns False when the two were already connected. That's exactly how you detect a cycle in an undirected graph: an edge between two things that are already connected closes a loop.

• Look for "connected", "groups", "same network" or "when does everything become connected".
• Mention both tricks, union by size and path compression, and that together they make it nearly O(1).
• Keep a group count as you go instead of counting at the end.

When does everyone become connected?

There are 6 people, numbered 0 to 5. You get logs of when two people became friends, sorted by time: (1: 0-1), (2: 3-4), (3: 2-3), (4: 1-5), (5: 2-4), (6: 0-3). When is everyone connected, directly or through others?

1. Start with 6 groups, one per person.
2. Time 1: join 0 and 1. 5 groups. Time 2: join 3 and 4. 4 groups.
3. Time 3: join 2 and 3. 2's group and 3's group differ, so 3 groups: {0, 1}, {2, 3, 4}, {5}.
4. Time 4: join 1 and 5. 2 groups: {0, 1, 5} and {2, 3, 4}.
5. Time 5: 2 and 4 are already in the same group, so union returns False and the count stays at 2.
6. Time 6: join 0 and 3. That merges the last two groups, so the count hits 1. The answer is 6.
7. Each log is a near-constant-time union, so it's O(m) for m logs, plus sorting if they're not sorted already. Running a BFS after each log would be O(m × (n + m)).

Things interviewers tend to ask once you've brought this up.

How do you count islands in a grid with union-find?

Give each land cell a number, union it with its land neighbors, and count the groups. It's handy when land gets added one cell at a time.

How do you find the redundant connection in a graph that should be a tree?

Union the edges in order. The first one where union returns False joins two things that were already connected, so it's the extra edge.

How do you group accounts that share an email?

Union each account's emails together. Accounts whose emails end up in the same group belong to the same person.

What's the cheapest way to connect every city?

Kruskal's algorithm: sort roads by cost and add each one that connects two different groups. Union-find tells you instantly whether they're already connected.

• Forgetting to call find on both items before comparing them.
• Skipping union by size and path compression, so chains get long.
• Counting groups at the end instead of keeping a running count.
• Using it when edges get removed, which it can't handle.
• Not mapping non-numeric items to numbers first.

• Graph traversal and shortest pathsBFS, DFS and Dijkstra, and when to use each.
• Trees and recursionSolve tree problems by deciding what each node returns.
• Hash maps and sliding windowsSwap nested loops for a single pass.