Grid Illumination LeetCode solution in Python

1
class Solution:
2
    def gridIllumination(
3
        self, n: int, lamps: List[List[int]], queries: List[List[int]]
4
    ) -> List[int]:
5
        rows = collections.Counter()
6
        cols = collections.Counter()
7
        diags1 = collections.Counter()
8
        diags2 = collections.Counter()
9
        lamps = {tuple(lamp) for lamp in lamps}
10

11
        for i, j in lamps:
12
            rows[i] += 1
13
            cols[j] += 1
14
            diags1[i + j] += 1
15
            diags2[i - j] += 1
16

17
        ans = []
18
        directions = (
19
            (-1, -1),
20
            (-1, 0),
21
            (-1, 1),
22
            (0, -1),
23
            (0, 0),
24
            (0, 1),
25
            (1, -1),
26
            (1, 0),
27
            (1, 1),
28
        )
29

30
        for i, j in queries:
31
            if rows[i] or cols[j] or diags1[i + j] or diags2[i - j]:
32
                ans.append(1)
33
            else:
34
                ans.append(0)
35

36
            for di, dj in directions:
37
                newI, newJ = i + di, j + dj
38
                if (newI, newJ) not in lamps:
39
                    continue
40
                lamps.remove((newI, newJ))
41
                rows[newI] -= 1
42
                cols[newJ] -= 1
43
                diags1[newI + newJ] -= 1
44
                diags2[newI - newJ] -= 1
45

46
        return ans