DSA Union-Find
Union-find is a structure that tracks groups of items and merges them. A parent array holds the groups as shallow trees.
First unions
Two unions merge two pairs of items. The parent array and two roots print below.
parent = list(range(6))
rank = [0] * 6
def find(x):
while parent[x] != x:
x = parent[x]
return x
def union(a, b):
ra, rb = find(a), find(b)
if ra == rb:
return
if rank[ra] < rank[rb]:
ra, rb = rb, ra
parent[rb] = ra
if rank[ra] == rank[rb]:
rank[ra] += 1
union(0, 1)
union(2, 3)
print(parent)
print(find(1), find(3))Path compression flattens a chain
Item 0 chains through 1 and 2 up to root 3. One find rewires every link on that path.
parent = [1, 2, 3, 3]
def find(x):
if parent[x] == x:
return x
parent[x] = find(parent[x])
return parent[x]
print("Before:", parent)
print("Root of 0:", find(0))
print("After:", parent)Example
class UnionFind:
def __init__(self, n):
self.parent = list(range(n))
self.rank = [0] * n
def find(self, x):
if self.parent[x] != x:
self.parent[x] = self.find(self.parent[x])
return self.parent[x]
def union(self, a, b):
ra, rb = self.find(a), self.find(b)
if ra == rb:
return False
if self.rank[ra] < self.rank[rb]:
ra, rb = rb, ra
self.parent[rb] = ra
if self.rank[ra] == self.rank[rb]:
self.rank[ra] += 1
return True
def depth(self, x):
d = 0
while self.parent[x] != x:
x = self.parent[x]
d += 1
return d
uf = UnionFind(10)
pairs = [(0, 1), (1, 2), (3, 4), (4, 5), (2, 3), (6, 7), (8, 9), (7, 8)]
for a, b in pairs:
uf.union(a, b)
print("Parents:", uf.parent)
print("Ranks:", uf.rank)
depths = [uf.depth(i) for i in range(10)]
print("Depths:", depths)
print("Max depth:", max(depths))
print("Roots:", sorted({uf.find(i) for i in range(10)}))
uf.find(5)
print("After find(5):", uf.parent)Key points
- Set: one group of items that belong together
- Parent array: a list where each item names its parent in its tree
- Root: the item whose parent is itself, naming its whole set
- Find: follow parent links from an item up to its root
- Union: attach one set's root under another set's root
- Path compression: a find rewires each item on the path to point at the root
- Union by rank: the shorter tree goes under the taller one
- Rank: a stored upper bound on a tree's height
In this example, a union-find over ten items runs eight unions. Two sets survive, named by roots 0 and 6. The depth list shows no item deeper than two steps from its root. A find on item 5 then compresses its path, and the parent array prints flatter. Union by rank kept the trees shallow, and compression flattens them further with use.
Note
Ranks change only when two trees of equal rank merge.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
What does path compression change?