Union-Find in Python: Disjoint Set Union with Path Compression
Who Is Connected to Whom?
A social network wants to know whether two people are linked through a chain of friends. An electricity company wants to know which towns are connected to the grid. A game wants to know whether two cells of a maze are joined. In each case, connections keep being added, and the question "are these two in the same group?" is asked again and again. The union-find data structure, also called the disjoint set union (DSU), answers it almost instantly, however many items there are.
This tutorial covers union-find in Python: the idea of disjoint sets, the two operations find and union, the two optimisations that make it fast (path compression and union by rank), and classic applications: counting connected components, detecting cycles in a graph, grouping accounts and counting islands. It is a key building block for Kruskal's minimum spanning tree algorithm. It builds on the tutorials on graphs and Big O notation.
The Idea: Sets with Representatives
Union-find keeps a collection of items divided into disjoint sets: groups that do not overlap. Each set is stored as a tree, and the item at the root of the tree is the set's representative. Two operations are supported:
- find(x): return the representative of the set containing x. Two items are in the same set exactly when they have the same representative.
- union(x, y): merge the sets containing x and y into one, by making one root point to the other.
Each item stores only its parent. A root is its own parent. Initially every item is in its own set.
A Simple Version
class SimpleDSU:
def __init__(self, n):
self.parent = list(range(n)) # everyone starts as their own root
def find(self, x):
while self.parent[x] != x: # climb to the root
x = self.parent[x]
return x
def union(self, a, b):
root_a, root_b = self.find(a), self.find(b)
if root_a != root_b:
self.parent[root_b] = root_a
d = SimpleDSU(6)
d.union(0, 1)
d.union(1, 2)
d.union(3, 4)
print(d.find(2) == d.find(0), d.find(3) == d.find(0))
print(d.parent)
Output:
True False [0, 0, 0, 3, 3, 5]
There are now three sets: {0, 1, 2}, {3, 4} and {5}. This version is correct, but in the worst case the trees become long chains, and find then takes O(n) time.
Making It Fast
1. Path compression
During find, once the root is known, make every item visited on the way point directly at the root. Future finds on those items then take a single step.
2. Union by rank (or by size)
When merging, attach the shorter tree under the root of the taller one, so that trees stay shallow. The rank is an upper bound on a tree's height. Tracking the size of each set works equally well, and it is useful in its own right.
With both optimisations, each operation takes amortised O(α(n)) time, where α is the inverse Ackermann function. It grows so slowly that α(n) is at most 4 for any n that could ever fit in a computer, so in practice each operation is effectively constant time.
class DSU:
def __init__(self, n):
self.parent = list(range(n))
self.rank = [0] * 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 self.parent[x] != root: # path compression
self.parent[x], x = root, self.parent[x]
return root
def union(self, a, b):
ra, rb = self.find(a), self.find(b)
if ra == rb:
return False # already in the same set
if self.rank[ra] < self.rank[rb]: # union by rank
ra, rb = rb, ra
self.parent[rb] = ra
self.size[ra] += self.size[rb]
if self.rank[ra] == self.rank[rb]:
self.rank[ra] += 1
self.count -= 1
return True
def connected(self, a, b):
return self.find(a) == self.find(b)
def set_size(self, x):
return self.size[self.find(x)]
d = DSU(8)
for a, b in [(0, 1), (2, 3), (1, 3), (5, 6)]:
d.union(a, b)
print(d.connected(0, 2), d.connected(0, 5))
print(d.count, d.set_size(3), d.set_size(7))
Output:
True False 4 4 1
The four sets are {0, 1, 2, 3}, {5, 6}, {4} and {7}. The union method returns False when the two items were already connected, which is exactly what cycle detection needs.
The path-compression loop uses tuple assignment: self.parent[x], x = root, self.parent[x] evaluates the right-hand side first, so it points x at the root and then moves on to x's old parent.
Application 1: Counting Connected Components
Given the roads between towns, how many separate road networks are there?
towns = ["Lagos", "Ibadan", "Abeokuta", "Enugu", "Onitsha", "Jos"]
index = {t: i for i, t in enumerate(towns)}
roads = [("Lagos", "Ibadan"), ("Ibadan", "Abeokuta"), ("Enugu", "Onitsha")]
net = DSU(len(towns))
for a, b in roads:
net.union(index[a], index[b])
groups = {}
for t in towns:
groups.setdefault(net.find(index[t]), []).append(t)
print(net.count, "networks:", list(groups.values()))
Output:
3 networks: [['Lagos', 'Ibadan', 'Abeokuta'], ['Enugu', 'Onitsha'], ['Jos']]
Breadth-first or depth-first search can also count components, but union-find shines when edges arrive one at a time and the question is asked after each one, because nothing has to be recomputed.
Application 2: Detecting a Cycle
In an undirected graph, adding an edge between two vertices that are already connected creates a cycle.
def first_cycle_edge(n, edges):
d = DSU(n)
for a, b in edges:
if not d.union(a, b):
return (a, b)
return None
print(first_cycle_edge(4, [(0, 1), (1, 2), (2, 3)]))
print(first_cycle_edge(4, [(0, 1), (1, 2), (2, 0), (2, 3)]))
Output:
None (2, 0)
This is the heart of Kruskal's algorithm for minimum spanning trees: take edges from cheapest to most expensive, and skip any edge that would form a cycle.
Application 3: Merging Accounts by Shared Email
Union-find works with any hashable items, not just numbers, if the parent links are kept in a dictionary.
class DictDSU:
def __init__(self):
self.parent = {}
def find(self, x):
self.parent.setdefault(x, x)
if self.parent[x] != x:
self.parent[x] = self.find(self.parent[x]) # recursive compression
return self.parent[x]
def union(self, a, b):
self.parent[self.find(a)] = self.find(b)
accounts = [
["ada@mail.com", "ada.o@school.edu"],
["bayo@mail.com"],
["ada.o@school.edu", "ada.work@firm.ng"],
]
d = DictDSU()
for emails in accounts:
for e in emails:
d.union(emails[0], e)
people = {}
for emails in accounts:
for e in emails:
people.setdefault(d.find(e), set()).add(e)
print(sorted(sorted(p) for p in people.values()))
Output:
[['ada.o@school.edu', 'ada.work@firm.ng', 'ada@mail.com'], ['bayo@mail.com']]
The first and third accounts share an email address, so they belong to the same person. The recursive find is short and clear, but for very deep trees an iterative version (as in the DSU class) avoids Python's recursion limit.
Application 4: Counting Islands
grid = [
"11000",
"11010",
"00011",
"10000",
]
rows, cols = len(grid), len(grid[0])
d = DSU(rows * cols)
land = 0
for r in range(rows):
for c in range(cols):
if grid[r][c] == "1":
land += 1
for dr, dc in ((1, 0), (0, 1)): # join with the cell below and to the right
nr, nc = r + dr, c + dc
if nr < rows and nc < cols and grid[nr][nc] == "1":
if d.union(r * cols + c, nr * cols + nc):
land -= 1
print("islands:", land)
Output:
islands: 3
Each land cell starts as its own island, and every successful union of two neighbouring land cells reduces the count by one. The cell at row r, column c is numbered r × cols + c.
Common Mistakes with Union-Find
- Comparing parents instead of roots.
parent[a] == parent[b]is not the test. Usefind(a) == find(b). - Linking the items instead of their roots in
union, as inparent[b] = a. This can split a tree apart. Always link root to root. - Skipping the optimisations on large inputs, which lets trees grow into long chains.
- Using union-find on directed graphs. It models undirected connectivity. Cycle detection in directed graphs needs depth-first search or topological sorting.
- Expecting to split sets. Union-find can merge sets but cannot efficiently undo a union.
Union-Find: Practice Exercises
- Given 10 students and pairs of friends, find the size of the largest friendship group.
- Given the edges added to a graph one by one, report the number of components after each edge.
- A set of equations such as "a==b", "b==c" and "a!=c" is given. Use union-find to decide whether they can all be true at once.
Sample solutions:
def largest_group(n, pairs):
d = DSU(n)
for a, b in pairs:
d.union(a, b)
return max(d.set_size(i) for i in range(n))
def components_after_each(n, edges):
d = DSU(n)
out = []
for a, b in edges:
d.union(a, b)
out.append(d.count)
return out
def equations_possible(equations):
d = DictDSU()
for eq in equations:
if eq[1:3] == "==":
d.union(eq[0], eq[3])
return all(d.find(eq[0]) != d.find(eq[3]) for eq in equations if eq[1:3] == "!=")
print(largest_group(10, [(0, 1), (1, 2), (3, 4), (5, 6), (6, 7), (7, 8)]))
print(components_after_each(5, [(0, 1), (2, 3), (1, 0), (3, 4), (0, 4)]))
print(equations_possible(["a==b", "b==c", "a!=c"]), equations_possible(["a==b", "c!=d"]))
Output:
4 [4, 3, 3, 2, 1] False True
Conclusion
Union-find keeps items in disjoint sets, each a tree whose root is the representative. find follows parent links to the root, and union links one root under another. With path compression and union by rank or size, each operation runs in effectively constant time. Use it to count connected components as edges arrive, to detect cycles in undirected graphs, to merge records that share keys and to count islands, and as the engine of Kruskal's minimum spanning tree algorithm.