Union-Find in Python
Union-Find (Disjoint Set Union, DSU) tracks 'which group does each element belong to' and answers `find` / `union` in nearly O(1) amortized when you implement both path compression and union by rank. It is the right data structure for connected-components, Kruskal's MST, and 'is X equivalent to Y' under a stream of merges. This entry covers a basic class, the rank + path compression upgrades, and an MST application.
550 views
4
class UnionFindNaive:
"""Each element points at its parent. Root of a tree is its own parent."""
def __init__(self, items):
self.parent = {x: x for x in items}
def find(self, x):
while self.parent[x] != x:
x = self.parent[x]
return x
def union(self, a, b):
ra, rb = self.find(a), self.find(b)
if ra == rb:
return False # already in same set
self.parent[ra] = rb
return True
def connected(self, a, b):
return self.find(a) == self.find(b)
uf = UnionFindNaive(range(6))
uf.union(0, 1)
uf.union(1, 2)
uf.union(3, 4)
print(uf.connected(0, 2)) # True (chain 0-1-2)
print(uf.connected(0, 3)) # False (different component)
print(uf.connected(3, 4)) # True
print(uf.connected(5, 5)) # True (singleton)
# Re-union returns False because the elements are already merged.
print(uf.union(0, 2)) # FalseThe simplest union-find is a parent pointer per element. find walks the parent chain to the root; union merges by pointing one root at the other. connected(a, b) is just 'do they share a root?'. Without optimizations, a worst-case sequence of unions builds a linear chain and find becomes O(n), so this naive form is fine for teaching but not for production. Returning False from union when the elements are already in the same set is a useful signal for cycle-detection problems like Kruskal's algorithm.
2 more snippets in this entry are available for premium members.
Upgrade to Premium