A disjoint set union is an ADT for managing an equivalence relation
over a collection of elements. Given a disjoint set dsu
, we
add the relationship x ~ y
by dsu.union(x, y)
and we check x ~ y
by
dsu.find(x) == dsu.find(y)
.
A simple DSU can be implemented with a list of integers if the size is fixed.
class DisjointSetUnion:
"""Disjoint set union implementation with a fixed array of ints"""
def __init__(self, size: int):
self.parents = [-1] * size
def find(self, key: int):
if self.parents[key] == -1:
return key
self.parents[key] = self.find(self.parents[key])
return self.parents[key]
def union(self, key1: int, key2: int):
= self.find(key1)
root1 = self.find(key2)
root2 if root1 == root2:
return
self.parents[root1] = root2
Another implementation can be done with a dictionary to allow string keys and a variable set size.
class DisjointSetUnion:
"""Disjoint set union implementation with a dictionary"""
def __init__(self):
self.parents = {}
def find(self, key: str):
if key not in self.parents:
return key
self.parents[key] = self.find(self.parents[key])
return self.parents[key]
def union(self, key1: str, key2: str):
= self.find(key1)
root1 = self.find(key2)
root2 if root1 == root2:
return
self.parents[root1] = root2