DSU with Path Compression
Implement find with path compression so that all nodes on the path point directly to the root, achieving near-O(1) amortised find.
What Is Disjoint Set Union?
Disjoint Set Union (DSU), also called Union-Find, is a data structure that maintains a collection of disjoint (non-overlapping) sets. It supports two core operations: find (which set does element x belong to?) and union (merge the sets containing x and y). DSU is ideal for dynamic connectivity problems where groups merge over time but never split.
Every element starts as its own set. As we process edges or relationships, we merge sets together. The challenge is doing this efficiently — naive implementations are O(n) per operation, but with optimisations we approach O(1) amortised.
# Naive DSU without optimisations
class DSU:
def __init__(self, n):
self.parent = list(range(n)) # each node is its own parent
def find(self, x):
while self.parent[x] != x:
x = self.parent[x]
return x
def union(self, x, y):
px, py = self.find(x), self.find(y)
if px != py:
self.parent[px] = pyThe Problem with Naive Find
In the naive DSU, find(x) walks up the parent chain until it reaches a node that points to itself (the root). If the tree is balanced, this is O(log n). But if we always union by linking the second root under the first, we can create a chain (degenerate tree) of length n, making each find O(n).
Consider unioning 0→1→2→3→4 in sequence. Node 0's find call must traverse the entire chain. With path compression, we eliminate this problem by making every visited node point directly to the root during the find operation itself.
# Worst case without compression: a chain
# parent = [1, 2, 3, 4, 4] => find(0) takes 4 steps
# After path compression: parent = [4, 4, 4, 4, 4] => find(0) takes 1 step
parent = [1, 2, 3, 4, 4]
print('Before:', parent)
# Simulate find(0) with naive approach
x = 0
steps = 0
while parent[x] != x:
x = parent[x]
steps += 1
print('Root:', x, 'Steps taken:', steps)