Bricks Falling When Hit LeetCode solution in Python

1
# Runtime: 7285 ms (Top 5.13%) | Memory: 29.8 MB (Top 5.13%)
2
from collections import defaultdict
3

4

5
class Solution:
6
    def hitBricks(self, grid: List[List[int]], hits: List[List[int]]) -> List[int]:
7
        parent = defaultdict()
8
        sz = defaultdict(lambda: 1)
9
        empty = set()
10

11
        def find(i):
12
            if parent[i] != i:
13
                parent[i] = find(parent[i])
14
            return parent[i]
15

16
        def union(i, j):
17
            pi = find(i)
18
            pj = find(j)
19
            if pi != pj:
20
                parent[pi] = pj
21
                sz[pj] += sz[pi]
22

23
        row = len(grid)
24
        col = len(grid[0])
25
        for r in range(row):
26
            for c in range(col):
27
                parent[(r, c)] = (r, c)
28
        parent[(row, col)] = (row, col)
29
        for r, c in hits:
30
            if grid[r][c]:
31
                grid[r][c] = 0
32
            else:
33
                empty.add((r, c))
34
        for r in range(row):
35
            for c in range(col):
36
                if not grid[r][c]:
37
                    continue
38
                for dr, dc in [[-1, 0], [1, 0], [0, 1], [0, -1]]:
39
                    if 0 <= r + dr < row and 0 <= c + dc < col and grid[r + dr][c + dc]:
40
                        union((r, c), (r + dr, c + dc))
41
                if r == 0:
42
                    union((r, c), (row, col))
43
        res = [0] * len(hits)
44
        for i in range(len(hits) - 1, -1, -1):
45
            r, c = hits[i]
46
            if (r, c) in empty:
47
                continue
48
            grid[r][c] = 1
49
            curbricks = sz[find((row, col))]
50
            for dr, dc in [[-1, 0], [1, 0], [0, 1], [0, -1]]:
51
                if 0 <= r + dr < row and 0 <= c + dc < col and grid[r + dr][c + dc]:
52
                    union((r, c), (r + dr, c + dc))
53
            if r == 0:
54
                union((r, c), (row, col))
55
            nextbricks = sz[find((row, col))]
56
            if nextbricks > curbricks:
57
                res[i] = nextbricks - curbricks - 1
58
        return res