Disjoint Set Forest

disjoint_set_forest.py

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)"

To test:

df = DisjointForest()
for i in range(ord('a'), ord('i')):
df.make_set(chr(i))
print(df.subsets)

df.union('a', 'c')
df.union('d', 'e')
df.union('e', 'f')
print(df.subsets)

df.union('a', 'd')
df.find('d')
df.find('e')
df.find('f')
print(df.subsets)

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)}