Solution
from typing import List
from itertools import pairwise
class UnionFind:
def __init__(self, n: int):
self.p = list(range(n))
self.size = [1] * n
def find(self, x: int):
if self.p[x] != x:
self.p[x] = self.find(self.p[x])
return self.p[x]
def union(self, a: int, b: int) -> bool:
pa, pb = self.find(a), self.find(b)
if pa == pb:
return False
if self.size[pa] > self.size[pb]:
self.p[pb] = pa
self.size[pa] += self.size[pb]
else:
self.p[pa] = pb
self.size[pb] += self.size[pa]
return True
class Solution:
def numIslands2(self, m: int, n: int, positions: List[List[int]]) -> List[int]:
uf = UnionFind(m * n)
grid = [[0] * n for _ in range(m)]
ans = []
dirs = (-1, 0, 1, 0, -1)
cnt = 0
for i, j in positions:
if grid[i][j]:
ans.append(cnt)
continue
grid[i][j] = 1
cnt += 1
for a, b in pairwise(dirs):
x, y = i + a, j + b
if 0 <= x < m and 0 <= y < n and grid[x][y]:
if uf.union(i * n + j, x * n + y):
cnt -= 1
ans.append(cnt)
return ans
if __name__ == "__main__":
sol = Solution()
print(sol.numIslands2(3,3,[[0,0],[0,1],[1,2],[2,1]]))
# [1,1,2,3]