Union-Find Data Structure for Efficient Connectivity Queries in Python
The Union-Find (or Disjoint Set Union) data structure efficiently manages a collection of disjoint sets, suppporting two core operations: determining which set an element belongs to (find) and merging two sets (union). It is particularly useful for modeling connectivity in graphs, networks, and dynamic equivalence relations.
A basic implementation uses two arrays: parent, where each node initially points to itself, and rank, which helps keep the tree shallow during unions.
class UnionFind:
def __init__(self, size):
self.parent = list(range(size))
self.rank = [0] * size
The find operation locates the root representative of an element’s set. Path compression optimizes future queries by flattening the structure as it traverses upward:
def find(self, x):
if self.parent[x] != x:
self.parent[x] = self.find(self.parent[x])
return self.parent[x]
The union operation merges two sets by attaching the root of the lower-rank tree under the root of the higher-rank one. If ranks are equal, one is chosen arbitrarily as the parent, and its rank is incremented:
def union(self, x, y):
rx, ry = self.find(x), self.find(y)
if rx == ry:
return
if self.rank[rx] < self.rank[ry]:
rx, ry = ry, rx
self.parent[ry] = rx
if self.rank[rx] == self.rank[ry]:
self.rank[rx] += 1
Connectivity between two elements is determined by comparing their roots:
def connected(self, x, y):
return self.find(x) == self.find(y)
The number of distinct connected components can be computed by counting unique root representatives:
def component_count(self):
roots = {self.find(i) for i in range(len(self.parent))}
return len(roots)
This structure is instrumental in algorithms like Kruskal’s minimum spanning tree, where it prevents cycles by checking if edge endpoints are already connected:
def kruskal_mst(edge_list, num_vertices):
edge_list.sort(key=lambda e: e[2])
dsu = UnionFind(num_vertices)
mst_edges = []
total_cost = 0
for u, v, weight in edge_list:
if not dsu.connected(u, v):
dsu.union(u, v)
mst_edges.append((u, v, weight))
total_cost += weight
return mst_edges, total_cost
edges = [(0, 1, 1), (0, 2, 4), (1, 2, 2), (1, 3, 5), (2, 3, 3)]
mst, cost = kruskal_mst(edges, 4)
Similarly, it models dynamic network connectivity:
class DynamicNetwork:
def __init__(self, n):
self.dsu = UnionFind(n)
def link(self, a, b):
self.dsu.union(a, b)
def reachable(self, a, b):
return self.dsu.connected(a, b)
net = DynamicNetwork(5)
net.link(0, 1)
net.link(1, 2)
assert net.reachable(0, 2)
assert not net.reachable(0, 4)
With path compresion and union by rank, both find and union operations run in nearly constant amortized time (inverse Ackermann complexity). While limited to equivalence and connectivity problems—and not well-suited for frequent element deletions—the Union-Find structure remains a compact, powerful tool for dynamic connectivity management.