947-uf.py

class UnionFind:
    def __init__(self):
        self.id = {}
        self.size = {}

    def setdefault(self, p: int):
        self.id.setdefault(p, p)
        self.size.setdefault(p, 1)
        
    def root(self, i: int) -> int:
        while i != self.id[i]:
            # path compression
            self.id[i] = self.id[self.id[i]]
            i = self.id[i]
        return i

    def connected(self, p: int, q: int) -> bool:
        return self.root(p) == self.root(q)

    def union(self, p: int, q: int) -> None:
        self.setdefault(p)
        self.setdefault(q)
        
        i = self.root(p)
        j = self.root(q)

        if i == j:
            return

        # make i always the smallest subtree
        if self.size[i] > self.size[j]:
            i, j = j, i

        # weighting
        self.id[i] = j
        self.size[j] += self.size[i]

class Solution:
    def removeStones(self, stones: List[List[int]]) -> int:
        if not stones: return 0
        
        uf = UnionFind()
        for i, j in stones:
            rowId, columnId = (i + 1), -(j + 1)
            uf.union(rowId, columnId)
            
        roots = [uf.root(i) for i in uf.id.values()]
        
        return len(stones) - len(set(roots))