Use this to learn the idea, then write your own version.
123 4class UnionFind(object):5 def __init__(self, n):6 self.set = range(n+1)7 self.size = [1]*(n+1)8 self.size[-1] = 09 10 def find_set(self, x):11 if self.set[x] != x:12 self.set[x] = self.find_set(self.set[x]) 13 return self.set[x]14 15 def union_set(self, x, y):16 x_root, y_root = map(self.find_set, (x, y))17 if x_root == y_root:18 return False19 self.set[min(x_root, y_root)] = max(x_root, y_root)20 self.size[max(x_root, y_root)] += self.size[min(x_root, y_root)]21 return True22 23 def top(self):24 return self.size[self.find_set(len(self.size)-1)]25 26 27class Solution(object):28 def hitBricks(self, grid, hits):29 """30 :type grid: List[List[int]]31 :type hits: List[List[int]]32 :rtype: List[int]33 """34 def index(C, r, c):35 return r*C+c36 37 directions = [(0, -1), (0, 1), (-1, 0), (1, 0)]38 R, C = len(grid), len(grid[0])39 40 hit_grid = [row[:] for row in grid]41 for i, j in hits:42 hit_grid[i][j] = 043 44 union_find = UnionFind(R*C)45 for r, row in enumerate(hit_grid):46 for c, val in enumerate(row):47 if not val:48 continue49 if r == 0:50 union_find.union_set(index(C, r, c), R*C)51 if r and hit_grid[r-1][c]:52 union_find.union_set(index(C, r, c), index(C, r-1, c))53 if c and hit_grid[r][c-1]:54 union_find.union_set(index(C, r, c), index(C, r, c-1))55 56 result = []57 for r, c in reversed(hits):58 prev_roof = union_find.top()59 if grid[r][c] == 0:60 result.append(0)61 continue62 for d in directions:63 nr, nc = (r+d[0], c+d[1])64 if 0 <= nr < R and 0 <= nc < C and hit_grid[nr][nc]:65 union_find.union_set(index(C, r, c), index(C, nr, nc))66 if r == 0:67 union_find.union_set(index(C, r, c), R*C)68 hit_grid[r][c] = 169 result.append(max(0, union_find.top()-prev_roof-1))70 return result[::-1]71 72