DisjointSet - union & find

disjoint_set.py

class DisjointSet(object):
    
    index = list()
    depth = list()
    
    def __init__(self, n):
        self.index = list(range(n))
        self.depth = [0] * n
        
    def tree_num(self):
        num = 0
        for i in range(len(self.index)):
            if self.index[i] == i:
                num += 1
        return num
    
    def find(self, i):
        if (self.index[i] != i):
            self.index[i] = self.find(self.index[i])    # optimize path
        return self.index[i]
    
    def union(self, i, j):
        index_i = self.find(i)
        index_j = self.find(j)
        if index_i != index_j:
            if self.depth[i] > self.depth[j]:
                self.index[index_j] = index_i
            elif self.depth[i] < self.depth[j]:
                self.index[index_i] = index_j
            else:
                self.index[index_i] = index_j
                self.depth[index_j] += 1

    def __str__(self):
        return 'items: {}\nindex: {}\ndepth: {}'.format(list(range(len(self.index))), self.index, self.depth)

test.py

%%time

js = DisjointSet(10)
links = [[1,2], [2,3], [3,5], [5, 6]]
for pair in links:
    js.union(pair[0], pair[1])
print(js.tree_num())
print(js)

# 6
# items: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
# index: [0, 2, 5, 2, 4, 5, 5, 7, 8, 9]
# depth: [0, 0, 1, 0, 0, 1, 0, 0, 0, 0]
# CPU times: user 4 ms, sys: 0 ns, total: 4 ms
# Wall time: 143 µs