LeetCode 947: Most Stones Removed with Same Row or Column

link

The relation “shares-row-or-column-with” is an equivalence relation on the set of stones. In other words, there is an undirected graph where stones are the vertices and an edge between two stones exist if they share a row or a column. In this stone graph, we can remove all but one stones from each connected component. We can use Union-Find to count connected components.

Note, the edges are implicit. For a given stone, to be able to find another stone sharing same row or same column (and thus forming an edge), we keep one stone per row and per col.

Say there are m rows, n columns, and s stones.

Time: \mathcal{O}(s \cdot \lg^*{s}), space: \mathcal{O}(s).

class UnionFind:
    def __init__(self, stones):
        self.parent_of, self.rank_of = {}, {}
        self.component_count = 0
        for s in stones:
            u = tuple(s)
            self.parent_of[u] = u
            self.rank_of[u] = 0
            self.component_count += 1

    def __compress_path(self, u, root):
        while u != (p := self.parent_of[u]):
            self.parent_of[u] = root
            u = p

    def __find(self, u):
        root = u
        while root != (p := self.parent_of[root]):
            root = p
        self.__compress_path(u, root)

        return root

    def connect(self, u, v):
        root_u = self.__find(u)
        root_v = self.__find(v)
        if root_u == root_v:
            return

        self.component_count -= 1

        rank_u = self.rank_of[root_u]
        rank_v = self.rank_of[root_v]
        if rank_u == rank_v:
            self.parent_of[root_u] = root_v
            self.rank_of[root_v] += 1
            return

        if rank_u > rank_v:
            self.parent_of[root_v] = root_u
        else:
            self.parent_of[root_u] = root_v


class Solution:
    def removeStones(self, stones: List[List[int]]) -> int:
        uf = UnionFind(stones)
        row_leads, col_leads = {}, {}
        for s in stones:
            u = tuple(s)
            r, c = u
            if v := row_leads.get(r, None):
                uf.connect(u, v)

            if v := col_leads.get(c, None):
                uf.connect(u, v)

            row_leads[r] = u
            col_leads[c] = u

        return len(stones) - uf.component_count

Leave a comment