0547 Number of Provinces

0547 Number of Provinces.py

class Solution:
    def findCircleNum(self, edges: List[List[int]]) -> int:
        n = len(edges)
        root = [i for i in range(n)]
        rank = [1] * n
        
        def find(i):
            r = root[i]
            while r != root[r]:
                root[r] = root[root[r]]
                r = root[r]
            root[i] = r
            return r
        
        def union(x, y):
            rx, ry = find(x), find(y)
            if rx == ry:
                return False
            if rank[rx] >= rank[ry]:
                rank[rx] += rank[ry]
                root[ry] = rx
            else:
                rank[ry] += rank[rx]
                root[rx] = ry
            return True
        
        for i in range(n):
            for j in range(i + 1, n):
                if edges[i][j] == 1:
                    union(i, j)
        res = set()
        for i in root:
            res.add(find(i))
        return len(res)