Skip to content

Instantly share code, notes, and snippets.

@davidrft
Created December 8, 2019 19:11
Show Gist options
  • Save davidrft/87e4511607bec741bf62ced5fdae1d60 to your computer and use it in GitHub Desktop.
Save davidrft/87e4511607bec741bf62ced5fdae1d60 to your computer and use it in GitHub Desktop.
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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment