L30 · Graphs IV — Union-Find (DSU)

Series: LeetCode — From Basics to Interview-Ready · Session 30 / 65 · Phase 1 · Core patterns

Why this session exists

Union-find is the highest ratio of power to lines in the entire interview canon. Fifteen lines, near-constant amortised time, and it trivialises a whole family of problems that are awkward with DFS: incremental connectivity, redundant edges, account merging, Kruskal's MST, and every "are these two things in the same group" query.

The reason it deserves its own session rather than a footnote in the graph chapter is that it answers a different question. DFS answers connectivity for a graph you already have. DSU answers connectivity for a graph that is being built one edge at a time, and it answers each query as the edges arrive. Re-running DFS after every edge insertion is O(E·(V+E)). DSU is O(E·α(V)), where α is inverse Ackermann and is at most 4 for any input that fits in the observable universe.

Memorise it. This is one of maybe four things in this series I would tell you to learn as literal muscle memory, because it never varies and it is always right.

Blank-file warm-up

Five minutes, empty file, no notes. Write from memory:

find(x)      — with path compression
union(a, b)  — by rank, returns False if already joined

The two things that go wrong. First, writing find without compression — it still works but degrades to O(n) per call on adversarial input. Second, forgetting the if ra == rb: return False guard, which is what lets union double as a cycle detector and as the component counter's trigger.

If you cannot produce both in five minutes today, this is the session to repeat until you can.

Pattern anatomy

The shape of problem that summons DSU: a stream of "these two things are now related" facts, plus questions about grouping or about whether a fact was redundant. The relation must be an equivalence — reflexive, symmetric, transitive. Friendship, connectivity, "same account", "same equation class". If the relation is directional, DSU does not apply; that is a directed-graph problem.

The invariant: every element points, directly or transitively, at a single canonical representative for its set, and two elements are in the same set if and only if they resolve to the same representative. Path compression flattens the pointer chains without changing which representative anything resolves to, which is why it is always safe.

The template. This is the version to memorise:

class DSU:
    def __init__(self, n):
        self.parent = list(range(n))     # every node is its own root
        self.rank = [0] * n              # upper bound on tree height
        self.count = n                   # number of disjoint components
 
    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                 # already joined: this edge is redundant
        if self.rank[ra] < self.rank[rb]:
            ra, rb = rb, ra              # attach the shorter tree under the taller
        self.parent[rb] = ra
        if self.rank[ra] == self.rank[rb]:
            self.rank[ra] += 1           # height grows only on a tie
        self.count -= 1
        return True

Four details that carry the whole thing.

self.parent[x] = self.parent[self.parent[x]] is path halving — each node on the walk is re-pointed at its grandparent. It is iterative, so no recursion limit, and it gives the same asymptotic guarantee as full path compression while being shorter than the two-pass recursive version.

Union by rank, not by raw size of subtree. rank is an upper bound on the tree height, and it only increments when you merge two trees of equal rank. Without it, a chain of unions in adversarial order produces a linked list and find becomes O(n).

return False on an already-joined pair. This single line is a cycle detector for undirected graphs and is the entire solution to Redundant Connection.

self.count starts at n and decrements on every successful union. That gives you the number of connected components in O(1) at any moment, without a second pass.

The functional form, when you do not want a class:

parent = list(range(n))
 
def find(x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x
 
def union(a, b):
    ra, rb = find(a), find(b)
    if ra == rb:
        return False
    parent[rb] = ra
    return True

Drop the rank and you drop one of the two optimisations. With path compression alone the amortised bound is still O(log n) per operation, which is fine for interview constraints — but say out loud that you are trading it away, do not just omit it silently.

The cue

You are looking at a union-find problem when the statement contains one of these tells:

1. Edges arrive as a stream and you are asked about connectivity as you go. Redundant Connection, Number of Operations to Make Network Connected. DFS would have to re-run from scratch after each edge.
2. "Are these two in the same group" asked many times. Each query must be O(1)-ish, and DSU is the only structure that gives you that on a mutable partition.
3. Merging entities by shared attributes. Accounts Merge (shared email), Similar String Groups. The attribute is the bridge; you union everything sharing it.
4. "Count connected components" on a static graph. DSU works and DFS works — pick either, but DSU is fewer lines and reads more directly.
5. Detecting a cycle in an undirected graph. The union returning False is the detection. Note this is undirected only; directed cycles need the three-colour DFS from session 29.
6. Minimum spanning tree. Kruskal's is: sort edges by weight, union them in order, skip the ones that return False. DSU is the entire supporting structure.

The negative cue: DSU cannot un-merge. If the problem involves removing edges or splitting groups, DSU is the wrong tool — you would process the edges in reverse and turn deletions into insertions, or use a different structure entirely.

Guided solve

Number of Connected Components in an Undirected Graph. Given n nodes labelled 0 to n−1 and a list of undirected edges, return the number of connected components.

Name the structure first. Every node starts in its own component, so the count starts at n. Every edge that joins two different components reduces the count by one. Every edge that joins two nodes already in the same component changes nothing — it is a redundant edge inside an existing component.

That is the whole solution stated in English, and the code is a direct transcription.

def countComponents(n, edges):
    parent = list(range(n))
    rank = [0] * n
    count = n
 
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
 
    for a, b in edges:
        ra, rb = find(a), find(b)
        if ra == rb:
            continue                     # same component already: no change
        if rank[ra] < rank[rb]:
            ra, rb = rb, ra
        parent[rb] = ra
        if rank[ra] == rank[rb]:
            rank[ra] += 1
        count -= 1
 
    return count

Trace on n = 5, edges = [[0,1],[1,2],[3,4]]. Start count 5. Edge (0,1): different roots, merge, count 4. Edge (1,2): find(1) resolves to 0, find(2) is 2, different, merge, count 3. Edge (3,4): different, merge, count 2. Answer 2 — the components are \{0,1,2\} and \{3,4\}.

Now the part worth saying aloud in an interview. The DFS solution to this problem is equally correct and roughly the same length. So why choose DSU? Because of the follow-up. "Now the edges arrive one at a time and after each one I want the current component count." DSU answers that in O(α) per edge with no code change at all. DFS has to re-traverse the whole graph after each edge, which is O(E·(V+E)) total. Volunteering that follow-up before the interviewer asks it is worth more than the solution itself.

Solo timed

Fifteen minutes each, timer visible, no editorial until it fires.

• Redundant Connection — return the edge that, when removed, leaves a tree. Think about what union returning False tells you and why processing edges in input order gives the last such edge, which is what the problem asks for.
• Accounts Merge — the nodes are not given as integers. Before writing any DSU, decide what you are unioning: account indices, or emails? Both work; one is noticeably cleaner. Pick one and commit.

If both land early, do Number of Operations to Make Network Connected. It needs two facts from the DSU: the component count, and how many edges were redundant. Both are already available in the template.

Common failure modes

Calling find on raw values instead of roots inside union. Writing parent[b] = a instead of parent[find(b)] = find(a) corrupts the structure — you re-point a non-root node and orphan whatever hung below its actual root. This is the single most damaging DSU bug and it produces wrong answers silently.

Omitting path compression. The code still works. On a path-shaped input built by adversarial union order, find degrades to O(n) and the whole thing goes quadratic. TLE with no wrong answer to debug.

Decrementing the component count unconditionally. The decrement belongs after the ra == rb guard. Decrementing on every edge gives n - len(edges), which can go negative.

Comparing parent[a] == parent[b] to test same-set. Two nodes in the same set can have different immediate parents; only their roots match. Always compare find(a) == find(b).

Forgetting to map non-integer entities to indices. Emails, strings, coordinates — you need a dict from entity to index before DSU can touch them. Build it in one pass, then run DSU over indices.

Using DSU for a directed graph. Union-find models an undirected equivalence relation. Applying it to directed edges loses the direction entirely, so Course Schedule cannot be solved this way — a valid DAG with a diamond looks identical to a cycle.

Complexity

Time: O(E · α(V)) for E operations, where α is the inverse Ackermann function. This is the classic Tarjan bound for path compression combined with union by rank. Because α(n) ≤ 4 for every practically representable n, each operation is treated as effectively constant, and the total is described as near-linear.

The bound is amortised, not worst-case per operation. A single find immediately after a long chain of unions can still walk several levels; the guarantee is that the total across all operations is bounded, because each walk flattens the structure and pays for future calls.

With path compression only: amortised O(log n) per operation. With union by rank only: worst-case O(log n) per operation. With neither: O(n) worst case, and a deliberately ordered input makes it quadratic overall.

Space: O(V). Two arrays of length V — parent and rank — plus O(1) for the counter. If entities are strings, add O(V) for the entity-to-index dict. Note that the iterative find uses no stack, which is a real advantage over the recursive form on large inputs.

Spaced queue

Re-solve whatever is due before starting anything new today. Status ladder:

• cold — solved unaided, first attempt clean → next review in 60 days
• warm — solved but slowly or with a stumble → 21 days
• hint — needed a nudge to get started → 7 days
• failed — could not produce a working solution → 2 days, then 7 days

DSU is the one template in this series I would put on a shorter cycle than the ladder suggests. Even if today's blank-file attempt was clean, re-write find and union from memory once a week until it is genuinely automatic. It costs ninety seconds and it is the highest-return ninety seconds in the whole graph section.