Union-Find with Path Compression
28 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)Cables are added between offices one at a time. After each cable you may ask: are these two offices on the same network? Union-Find answers in almost constant time and never needs to search the graph again.
Groups that only merge
Union-Find, also called disjoint set union, keeps items in non-overlapping groups. find(x) returns the leader of x's group and union(a, b) merges the groups of a and b.
Start with every item alone in its own group. Each group keeps one item as its leader. Find returns the leader of an item. Union joins two groups by making one leader point to the other.
Example 1 · Brute force: relabel the whole group
n = 6label = list(range(n))
def union(a, b): old, new = label[b], label[a] for i in range(n): if label[i] == old: label[i] = new
union(0, 1)union(2, 3)union(1, 3)print(label)print(label[0] == label[2])[0, 0, 0, 0, 4, 5] True
What this code does
- Every union scans all n items, so m unions cost O(n * m). The check is O(1).
- It is fine for tiny inputs, but too slow when both n and m are 100000.
Find with path compression
To find a leader, follow parent links up. A long chain makes this slow. After the search, point every item on the chain straight at the leader, so the next search takes one step.
def find(parent, x): if parent[x] != x: parent[x] = find(parent, parent[x]) return parent[x]Example 2 · Path compression flattens a chain
def find(parent, x): if parent[x] != x: parent[x] = find(parent, parent[x]) return parent[x]
parent = [0, 0, 1, 2, 3, 4]print(parent)print(find(parent, 5))print(parent)[0, 0, 1, 2, 3, 4] 0 [0, 0, 0, 0, 0, 0]
What this code does
- The parent list is a chain 5 → 4 → 3 → 2 → 1 → 0. The first find walks all five links.
- After it, every item points at 0, so later finds take one step.
The recursion returns the leader and rewrites each parent on the way back. Recursive depth can reach n, so the class later in this lesson uses a loop instead.
Union by size
When two groups merge, hang the smaller tree under the leader of the bigger one, so trees stay short. With path compression, each operation costs about O(α(n)), at most 4 for real inputs.
def union(parent, size, a, b): ra, rb = find(parent, a), find(parent, b) if ra == rb: return False if size[ra] < size[rb]: ra, rb = rb, ra parent[rb] = ra size[ra] += size[rb] return TrueExample 3 · A DSU class
class DSU: def __init__(self, n): self.parent = list(range(n)) self.size = [1] * n self.count = n
def find(self, x): while self.parent[x] != x: self.parent[x] = self.parent[self.parent[x]] 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] self.count -= 1 return True
dsu = DSU(6)print(dsu.union(0, 1))print(dsu.union(2, 3))print(dsu.union(1, 3))print(dsu.union(0, 2))print(dsu.count, dsu.size[dsu.find(0)])True True True False 3 4
What this code does
- The fourth call returns False, because 0 and 2 are already in the same group.
- The loop in
findmoves each item to its grandparent. This halves the path and avoids recursion.
It returns False when a and b already share a leader. That answer is also how you notice that a new edge would close a cycle.
Components and cycles
Start with n groups. Every successful union reduces the count by one. After reading all edges, the count is the number of connected components. A union that returns False means the edge joins two items that were already connected.
Example 4 · Count connected components
class DSU: def __init__(self, n): self.parent = list(range(n)) self.size = [1] * n self.count = n
def find(self, x): while self.parent[x] != x: self.parent[x] = self.parent[self.parent[x]] 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] self.count -= 1 return True
n = 7edges = [(0, 1), (1, 2), (3, 4), (5, 5)]dsu = DSU(n)for a, b in edges: dsu.union(a, b)print(dsu.count)4
What this code does
- Groups are {0, 1, 2}, {3, 4}, {5} and {6}. The self-edge (5, 5) changes nothing.
- Time O((n + m) * α(n)), space O(n).
Example 5 · Find the edge that closes a cycle
class DSU: def __init__(self, n): self.parent = list(range(n)) self.size = [1] * n self.count = n
def find(self, x): while self.parent[x] != x: self.parent[x] = self.parent[self.parent[x]] 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] self.count -= 1 return True
edges = [(0, 1), (1, 2), (2, 3), (3, 1), (3, 4)]dsu = DSU(5)for a, b in edges: if not dsu.union(a, b): print('cycle closed by', (a, b)) breakcycle closed by (3, 1)
What this code does
- The edge (3, 1) is the first one whose ends are already connected through 1, 2 and 3.
- In an undirected graph a union that returns False always means a cycle.
Example 6 · List the groups
from collections import defaultdict
class DSU: def __init__(self, n): self.parent = list(range(n)) self.size = [1] * n self.count = n
def find(self, x): while self.parent[x] != x: self.parent[x] = self.parent[self.parent[x]] 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] self.count -= 1 return True
dsu = DSU(8)for a, b in [(0, 4), (4, 6), (1, 5), (2, 3), (3, 7)]: dsu.union(a, b)groups = defaultdict(list)for x in range(8): groups[dsu.find(x)].append(x)print(sorted(groups.values()))[[0, 4, 6], [1, 5], [2, 3, 7]]
What this code does
- Group items by their leader with a
defaultdict(list). Each leader is the key of one group. - Sorting the list of groups gives the same output on every run.
Next step
Check your understanding
11 quick questions