Skip to content

Instantly share code, notes, and snippets.

@mkolod
Last active January 7, 2021 05:24
Show Gist options
  • Save mkolod/87fafe9ccf8ee62378b6e0f939ae5604 to your computer and use it in GitHub Desktop.
Save mkolod/87fafe9ccf8ee62378b6e0f939ae5604 to your computer and use it in GitHub Desktop.
Disjoint Set Forest
class DisjointForest:
class Subset:
def __init__(self, elem, parent=None, rank=0):
self.elem = elem
self.parent = parent
self.rank = rank
def __repr__(self):
parent = "self" if self.parent == self else repr(self.parent)
return f"Subset(elem={self.elem}, parent={parent}, rank={self.rank})"
def __eq__(self, other):
return self.elem == other.elem
def __init__(self):
self.subsets = {}
def make_set(self, x):
temp = DisjointForest.Subset(x)
temp.parent = temp
self.subsets[x] = temp
def find(self, x):
if x not in self.subsets:
raise Exception(f"{x} is not a member of the disjoint forest")
temp = self.subsets[x]
if temp.parent == temp:
return temp
else:
# path compression
temp.parent = self.find(temp.parent.elem)
self.subsets[temp.parent.elem] = temp.parent
return temp.parent
def union(self, a, b):
if a not in self.subsets:
raise Exception(f"{x} is not a member of the disjoint forest")
if b not in self.subsets:
raise Exception(f"{x} is not a member of the disjoint forest")
a2 = self.find(a)
b2 = self.find(b)
if a2 == b2:
return
else:
if a2.rank < b2.rank:
# we merge smaller rank set to bigger rank set
# If ranks are equal, it doesn't matter which way we merge
a2, b2 = b2, a2
b2.parent = a2
# If ranks are equal, we increment x's rank by 1
if a2.rank == b2.rank:
a2.rank += 1
def __repr__(self):
return f"DisjointForest(\n{self.subsets}\n)"
@mkolod
Copy link
Author

mkolod commented Jan 7, 2021

Result of above test code:

{'a': Subset(elem=a, parent=self, rank=1),
 'b': Subset(elem=b, parent=Subset(elem=a, parent=self, rank=1), rank=0),
 'c': Subset(elem=c, parent=Subset(elem=a, parent=self, rank=1), rank=0),
 'd': Subset(elem=d, parent=self, rank=1),
 'e': Subset(elem=e, parent=Subset(elem=d, parent=self, rank=1), rank=0),
 'f': Subset(elem=f, parent=Subset(elem=d, parent=self, rank=1), rank=0),
 'g': Subset(elem=g, parent=self, rank=0),
 'h': Subset(elem=h, parent=self, rank=0)}


{'a': Subset(elem=a, parent=self, rank=2),
 'b': Subset(elem=b, parent=Subset(elem=a, parent=self, rank=2), rank=0),
 'c': Subset(elem=c, parent=Subset(elem=a, parent=self, rank=2), rank=0),
 'd': Subset(elem=d, parent=Subset(elem=a, parent=self, rank=2), rank=1),
 'e': Subset(elem=e, parent=Subset(elem=d, parent=Subset(elem=a, parent=self, rank=2), rank=1), rank=0),
 'f': Subset(elem=f, parent=Subset(elem=d, parent=Subset(elem=a, parent=self, rank=2), rank=1), rank=0),
 'g': Subset(elem=g, parent=self, rank=0),
 'h': Subset(elem=h, parent=self, rank=0)}

{'a': Subset(elem=a, parent=self, rank=2),
 'b': Subset(elem=b, parent=Subset(elem=a, parent=self, rank=2), rank=0),
 'c': Subset(elem=c, parent=Subset(elem=a, parent=self, rank=2), rank=0),
 'd': Subset(elem=d, parent=Subset(elem=a, parent=self, rank=2), rank=1),
 'e': Subset(elem=e, parent=Subset(elem=a, parent=self, rank=2), rank=0),
 'f': Subset(elem=f, parent=Subset(elem=a, parent=self, rank=2), rank=0),
 'g': Subset(elem=g, parent=self, rank=0),
 'h': Subset(elem=h, parent=self, rank=0)}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment