Union-Find in Python: what it is and when a problem needs it

Short answer. Union-Find is a structure that tracks which elements share a group. It does two things fast. Tell whether two elements are connected. Merge two groups. With two small tricks each operation is close to O(1). Below: the Python code, the wording that calls for it, and the costly mistakes.

What is Union-Find?

Every element points at a parent. Follow the parents up and you reach a root. Two elements are in the same group when they share a root.

That is the whole idea. Two operations sit on top of it.

• find(x) walks up from x and returns its root.
• union(a, b) finds both roots and makes one point at the other.

The other name is Disjoint Set Union, or DSU. Same structure, different label.

How do you write Union-Find in Python?

Elements are numbered 0 to n-1. Each starts as its own root.

class UnionFind:
    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]]  # path halving
            x = self.parent[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]
        return True

Two lines do the heavy lifting.

• Path halving in find. Every node on the way up skips a step. Trees stay flat.
• Union by size. The smaller tree hangs under the larger one, so trees stay shallow.

With both, a sequence of m operations costs O(m α(n)). α grows so slowly that it is at most 4 in practice.

union returns False when the two were already connected. That one bit is the answer to half the problems below.

When does a problem need Union-Find?

The wording again.

• "How many connected components" or "how many provinces".
• "Are these two nodes connected" asked many times.
• "Which edge closes a cycle" in an undirected graph.
• "Merge accounts", "group the emails", "friend circles".
• Edges arrive one at a time, with a question after each.

The common thread: groups that only merge, never split, and many membership queries.

Signals, near misses and six code templates are on the Union-Find pattern page.

How is Union-Find different from BFS or DFS?

Both answer "are these connected". They differ in when the graph is known.

	Union-Find	BFS / DFS
Graph is	Built edge by edge	Known in full up front
Best at	Many connectivity queries	One traversal, paths, distances
Finds the path?	No, only yes or no	Yes
Cost per query	Near O(1)	O(V + E)

Need the actual path or the shortest one? That is BFS or DFS. Need to count groups or catch the edge that closes a cycle? Union-Find.

Which problems use Union-Find?

Five that come up in interviews more than the rest.

• Number of Connected Components. Start with n groups, subtract one per successful union.
• Redundant Connection. The first union that returns False names the edge.
• Number of Provinces. The same count on an adjacency matrix.
• Accounts Merge. Union every email in one account, then group by root.
• Kruskal's minimum spanning tree. Sort edges, union, skip the ones that return False.

How do you solve Redundant Connection with Union-Find?

The problem. A tree with n nodes got one extra edge. Find an edge whose removal gives a tree again. Return the last one in the input if there are several.

A tree has no cycles. The extra edge is the one that closes a cycle.

So walk the edges in order and union each pair. The first union that returns False is the answer.

def redundant_connection(edges):
    uf = UnionFind(len(edges) + 1)
    for a, b in edges:
        if not uf.union(a, b):
            return [a, b]

Trace it on edges (1, 2), (1, 3), (2, 3). The first two unions succeed. The third finds 2 and 3 already share root 1. Answer: [2, 3].

Four lines, and the whole difficulty was in the reading. That is typical for this pattern.

How do you explain Union-Find in an interview?

Three sentences before the code.

• "The edges only ever join groups. I need fast membership and fast merge, so Union-Find."
• "Each node points at a parent. Same root means same group."
• "Path compression and union by size keep every operation near constant."

Then the class. Then the trace above. Name the two tricks again at the end; they are what separates a pass from a partial.

What goes wrong with Union-Find?

Four mistakes cover almost every failed submission.

• Comparing parents instead of roots. Two nodes can share a group and have different parents.
• Skipping path compression. Trees turn into chains and find becomes O(n).
• Using it on a directed graph. Union-Find has no direction; "A reaches B" is not "B reaches A".
• Rebuilding the structure per query. The point is to keep one and reuse it.

How do you make Union-Find stick?

Write the class from memory, both tricks included. Then solve Number of Connected Components and Redundant Connection back to back.

Recognising the pattern from the wording is the harder half. That is what LeetCode patterns is about. The placement test measures it.

FAQ

What is the time complexity of Union-Find?

Near O(1) per operation with path compression and union by size. Formally O(α(n)), where α is below 5 in practice.

Does Union-Find work on a directed graph?

No. It only tracks undirected connectivity. For direction you need a traversal.

Can Union-Find split a group?

No. Groups only merge. A problem that removes edges is usually solved backwards, adding them instead.

Union by size or union by rank?

Either. Both keep trees shallow. Size is easier to explain and also answers "how big is this group".