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 rows,
columns, and
stones.
Time: , space:
.
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