Union Find

Union-Find (Disjoint Set Union)

# Efficiently groups elements into disjoint sets
# find, union: near O(1) with path compression + union by rank

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank   = [0] * n
        self.count  = n  # number of components

    def find(self, x) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # path compression
        return self.parent[x]

    def union(self, x, y) -> bool:
        px, py = self.find(x), self.find(y)
        if px == py: return False  # already connected
        # Union by rank
        if self.rank[px] < self.rank[py]: px, py = py, px
        self.parent[py] = px
        if self.rank[px] == self.rank[py]: self.rank[px] += 1
        self.count -= 1
        return True

    def connected(self, x, y) -> bool:
        return self.find(x) == self.find(y)

# Number of provinces
def num_provinces(isConnected):
    n  = len(isConnected)
    uf = UnionFind(n)
    for i in range(n):
        for j in range(i+1, n):
            if isConnected[i][j]:
                uf.union(i, j)
    return uf.count