~/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.

what

A forest of parent pointers: each group is a tree, named by its root. Merge by linking roots; flatten paths as you walk them.

use when

Things get connected over time and you keep asking whether two are connected, or how many groups there are.

time

O(α(n))

space

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 x to 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.

  1. Start: parent[i] = i for every item. Every item is its own group of size 1, and its own root.
  2. 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.
  3. find(x): follow parent from x until you reach a node that is its own parent. That node is the root, and it names the group.
  4. Path compression: walk the same path again and point each node straight at the root. The next find from any of them takes one step.
  5. 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.
  6. Count groups by starting at n and subtracting one for every union that actually linked two roots.
loading union-find…
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 a 3 more than b?”.
  • 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. 1

    Friend pairs arrive one at a time; after each, you must report how many friend groups exist. Which structure fits best?

  2. 2

    What does this print?

    N = 6
    parent = list(range(N))
    def find(x):
    root = x
    while parent[root] != root: root = parent[root]
    while parent[x] != root: parent[x], x = root, parent[x]
    return root
    def union(a, b):
    ra, rb = find(a), find(b)
    if ra == rb: return False
    parent[rb] = ra
    return True
    count = N
    for a, b in [(0, 1), (1, 2), (0, 2), (3, 4)]:
    if union(a, b): count -= 1
    print(count, find(0) == find(2), find(2) == find(3))
  3. 3

    This union forgets to find the roots. What does it print?

    parent = list(range(3))
    def find(x):
    while parent[x] != x: x = parent[x]
    return x
    def union(a, b):
    parent[a] = b # bug
    union(0, 1)
    union(0, 2)
    print(find(1) == find(2))
  4. 4

    A solution keeps count = n and does union(a, b); count -= 1 for each edge. Where does it go wrong?

  5. 5

    With both path compression and union by rank (or size), what’s the amortized cost per find/union?

Practice problems

Further reading

esc