disjoint_set.py

class DisjointSet:
    def __init__(self, elems):
        self.elems = elems
        self.parent = {}
        self.size = {}
        for elem in elems:
            self.make_set(elem)

    def make_set(self, x):
        self.parent[x] = x
        self.size[x] = 1

    def find(self, x):
        if self.parent[x] == x:
            return x
        else:
            self.parent[x] = self.find(self.parent[x])
            return self.parent[x]

    def union(self, x, y):
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x == root_y:
            return
        elif self.size[root_x] < self.size[root_y]:
            self.parent[root_x] = root_y
            self.size[root_y] += self.size[root_x]
        else:
            self.parent[root_y] = root_x
            self.size[root_x] += self.size[root_y]

            
data = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']
ds = DisjointSet(data)
ds.union('B', 'F')
print(ds.parent)
ds.union('F', 'H')
print(ds.parent)
print(ds.find('H'))
ds.union('A', 'G')
print(ds.parent)
ds.union('F', 'A')
print(ds.parent)
print(ds.find('G'))
ds.union('C', 'D')
print(ds.parent)
ds.union('C', 'E')
print(ds.parent)
print(ds.find('E'))
ds.union('D', 'A')
print(ds.parent)