Algorithms

disjoint-set-union.py

class DSU:
    """
    Disjoint set union with union-by-rank
    """
    def __init__(self, n):
        # may need to alter the total number, like (n + 1)
        self.parent = list(range(n))
        self.rank = [0] * n
        # count of dsu
        self.count = n
    
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def union(self, x, y):
        i, j = self.find(x), self.find(y)
        if i != j:
            if self.rank[i] < self.rank[j]:
                self.parent[i] = j
            elif self.rank[i] > self.rank[j]:
                self.parent[j] = i
            else:
                self.parent[i] = j
                self.rank[j] += 1

            self.count -= 1