Group elements into sets and instantly check if two items belong to the same group.
A Disjoint Set Union (DSU) structure tracks a collection of elements partitioned into non-overlapping sets. It supports two operations almost instantly:
Find: Determine which set an element belongs to by following its parent pointer to a "root" representative. We use Path Compression to point all visited nodes directly to the root, making future lookups O(1).
Union: Merge two sets by pointing the root of one to the root of another. We use Union by Rank to attach the smaller tree under the taller tree, keeping it flat.
class DSU:
def __init__(self, n):
self.parent = list(range(n))
self.rank = [1] * n
def find(self, i):
if self.parent[i] != i:
# Path compression
self.parent[i] = self.find(self.parent[i])
return self.parent[i]
def union(self, i, j):
root_i = self.find(i)
root_j = self.find(j)
if root_i != root_j:
# Union by rank
if self.rank[root_i] < self.rank[root_j]:
self.parent[root_i] = root_j
elif self.rank[root_i] > self.rank[root_j]:
self.parent[root_j] = root_i
else:
self.parent[root_j] = root_i
self.rank[root_i] += 1