~/data-structures/union-find
Union-find
Track which items belong together as you merge groups. Both merging and asking "same group?" take almost constant time.
A forest of parent pointers: each group is a tree, named by its root. Merge by linking roots; flatten paths as you walk them.
Things get connected over time and you keep asking whether two are connected, or how many groups there are.
O(α(n))
O(n)
You’ll recognise it when
- Items get merged into groups over time: accounts that share an email, cities joined by new roads, pixels of the same colour.
- You keep asking “are these two in the same group?” or “how many groups are left?”
- You’re building a minimum spanning tree with Kruskal’s algorithm.
- An edge list arrives one edge at a time and you need to know when an edge closes a cycle.
If the graph is fixed and you only ask once, a single BFS or DFS is just as good. Union-find shines when merges and questions are mixed together.
The idea
Give every group a leader. Each item remembers someone in its group, its parent, and following parents always ends at the leader, the root, which is its own parent. To answer “same group?”, walk both items up to their roots and compare.
Merging two groups is a single pointer change: make one root the parent of the other. No items are moved or relabelled.
Two small tricks keep the walks short:
- Union by size. Hang the smaller tree under the bigger one, never the other way round. A node only gets deeper when its tree at least doubles in size, so no path is ever longer than log₂ n.
- Path compression. After walking from
xto its root, point every node on that path straight at the root. The next walk from any of them is one step.
How it works
Scroll through the steps and the graphic follows along. The example builds two trees of four, merges them into one tree three levels deep, then shows find(7) flattening it.
- Start:
parent[i] = ifor every item. Every item is its own group of size 1, and its own root. union(a, b): find the root of each. If the roots differ, make one root the parent of the other. The arrow is the whole merge: every node below comes along.find(x): followparentfromxuntil you reach a node that is its own parent. That node is theroot, and it names the group.- Path compression: walk the same path again and point each node straight at the
root. The nextfindfrom any of them takes one step. - Union by size: when linking, the root of the smaller tree goes under the root of the bigger one, even if it was named first. Add the sizes at the new root.
- Count groups by starting at
nand subtracting one for every union that actually linked two roots.
class DSU:
def __init__(self, n):
self.parent = list(range(n))
self.size = [1] * n
self.count = n # number of separate sets
def find(self, x):
root = x
while self.parent[root] != root:
root = self.parent[root]
while x != root: # path compression
nxt = self.parent[x]
self.parent[x] = root
x = nxt
return root
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
def connected(self, a, b):
return self.find(a) == self.find(b)#include <numeric>
#include <utility>
#include <vector>
using namespace std;
struct DSU {
vector<int> parent, size;
int count; // number of separate sets
DSU(int n) : parent(n), size(n, 1), count(n) {
iota(parent.begin(), parent.end(), 0);
}
int find(int x) {
int root = x;
while (parent[root] != root) {
root = parent[root];
}
while (x != root) { // path compression
int next = parent[x];
parent[x] = root;
x = next;
}
return root;
}
bool unite(int a, int b) {
int ra = find(a), rb = find(b);
if (ra == rb) return false;
if (size[ra] < size[rb]) swap(ra, rb);
parent[rb] = ra;
size[ra] += size[rb];
count--;
return true;
}
bool connected(int a, int b) { return find(a) == find(b); }
};class DSU {
int[] parent, size;
int count; // number of separate sets
DSU(int n) {
parent = new int[n];
size = new int[n];
count = n;
for (int i = 0; i < n; i++) {
parent[i] = i;
size[i] = 1;
}
}
int find(int x) {
int root = x;
while (parent[root] != root) {
root = parent[root];
}
while (x != root) { // path compression
int next = parent[x];
parent[x] = root;
x = next;
}
return root;
}
boolean union(int a, int b) {
int ra = find(a), rb = find(b);
if (ra == rb) return false;
if (size[ra] < size[rb]) {
int t = ra; ra = rb; rb = t;
}
parent[rb] = ra;
size[ra] += size[rb];
count--;
return true;
}
boolean connected(int a, int b) { return find(a) == find(b); }
}find is written with loops rather than recursion on purpose. A recursive find is shorter, but before compression kicks in a path can be long enough to blow Python’s recursion limit of 1,000.
Why it’s (almost) O(1)
With both tricks, any sequence of m operations on n items takes O(m · α(n)) time, where α is the inverse Ackermann function. α grows so slowly that it’s at most 4 for any n you could store in the universe, so treat each operation as constant time.
| Version | Cost per operation |
|---|---|
| No tricks | O(n) worst case: a chain |
| Union by size only | O(log n) |
| Path compression only | O(log n) amortised |
| Both | O(α(n)), effectively O(1) |
Space is two arrays of n integers.
Common mistakes
Linking the items instead of their roots
union(a, b) must connect the roots. Setting parent[a] = b detaches a from the rest of its own group.
self.parent[a] = b # ✗ a leaves its old group behind
self.parent[find(a)] = find(b) # ✓ the whole group moves
Comparing parents instead of roots
Two items are connected when their roots match. Their parents can differ even when they’re in the same group, especially before paths are compressed.
parent[a] == parent[b] # ✗ misses deeper members
find(a) == find(b) # ✓
Recursion depth in find
A recursive find works in C++ and Java for most inputs but crashes Python on a long chain. Use the two-loop version above, or raise the limit and hope.
Forgetting that union can fail
When both roots are equal, nothing is merged. If you count groups or detect cycles, only act when union returns True: in Kruskal’s algorithm, that edge would form a cycle and must be skipped.
Variations
- Union by rank stores an upper bound on each tree’s height instead of its size. The guarantee is the same; size is handier because you often want group sizes anyway.
- Randomised linking. Pick which root goes on top with a coin flip. It is simpler and, with path compression, fast in expectation.
- Extra data per group. Keep a sum, minimum or member list at each root and combine them in
union. - Weighted union-find stores an offset from each node to its parent, for questions like “is
a3 more thanb?”. - Map arbitrary keys (emails, coordinates) to integers first with a dictionary, then use the array version.
Check yourself
5 quick questions. Pick an answer to see why it's right or wrong.
-
1
Friend pairs arrive one at a time; after each, you must report how many friend groups exist. Which structure fits best?
Union-Find handles incremental merges in near-constant amortized time, and the count only changes when a union actually joins two different roots. Re-running BFS is correct but O(V + E) per query, which is O(E·(V + E)) overall.
-
2
What does this print?
N = 6parent = list(range(N))def find(x):root = xwhile parent[root] != root: root = parent[root]while parent[x] != root: parent[x], x = root, parent[x]return rootdef union(a, b):ra, rb = find(a), find(b)if ra == rb: return Falseparent[rb] = rareturn Truecount = Nfor a, b in [(0, 1), (1, 2), (0, 2), (3, 4)]:if union(a, b): count -= 1print(count, find(0) == find(2), find(2) == find(3))Starting with 6 singletons: (0,1), (1,2) and (3,4) succeed, and (0,2) fails because they’re already joined. That’s 3 successful unions, so 6 - 3 = 3 components: {0,1,2}, {3,4}, {5}. 0 and 2 are connected; 2 and 3 are not.
-
3
This
unionforgets to find the roots. What does it print?parent = list(range(3))def find(x):while parent[x] != x: x = parent[x]return xdef union(a, b):parent[a] = b # bugunion(0, 1)union(0, 2)print(find(1) == find(2))union(0, 1)setsparent[0] = 1. Thenunion(0, 2)overwritesparent[0] = 2, detaching 0 from 1. Now 1 is its own root and 2 is another, so 1 and 2 look disconnected. Always link the roots:parent[find(a)] = find(b). -
4
A solution keeps
count = nand doesunion(a, b); count -= 1for each edge. Where does it go wrong?Only a union that links two different roots reduces the component count. Have
unionreturn whether it merged and decrement only then. Graphs with cycles or duplicated input edges are exactly the test cases that catch this. -
5
With both path compression and union by rank (or size), what’s the amortized cost per
find/union?The combination gives O(α(n)) amortized, which is ≤ 4 for any realistic n. Union by rank alone gives O(log n) per op; neither guarantees O(1) worst case for a single operation, since one
findcan still walk a path before compressing it.