leetcode/grid-illumination.py at master · cannium/leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
def lit(qx, qy, diag, diag2, row, col):
    d = qx - qy
    if d in diag:
        return 1
    d2 = qy + qx
    if d2 in diag2:
        return 1
    if qx in row:
        return 1
    if qy in col:
        return 1
    return 0

direction = [
    (-1, -1),
    (-1, 0),
    (-1, 1),
    (0, -1),
    (0, 0),
    (0, 1),
    (1, -1),
    (1, 0),
    (1, 1)
]

def turnOff(qx, qy, board, N):
    ans = []
    for dx, dy in direction:
        x, y = qx + dx, qy + dy
        if x < 0 or x >= N:
            continue
        if y < 0 or y >= N:
            continue
        if (x, y) in board:
            ans.append((x,y))
    return ans


class Solution(object):
    def gridIllumination(self, N, lamps, queries):
        """
        :type N: int
        :type lamps: List[List[int]]
        :type queries: List[List[int]]
        :rtype: List[int]
        """
        board = {}
        diag = {}
        diag2 = {}
        row = {}
        col = {}
        for l in lamps:
            l = tuple(l)
            x, y = l
            board[l] = 1
            d = x - y
            if d in diag:
                diag[d][l] = 1
            else:
                diag[d] = {l: 1}
            d2 = y + x
            if d2 in diag2:
                diag2[d2][l] = 1
            else:
                diag2[d2] = {l: 1}
            if x in row:
                row[x][l] = 1
            else:
                row[x] = {l: 1}
            if y in col:
                col[y][l] = 1
            else:
                col[y] = {l: 1}
        ans = []
        for qx, qy in queries:
            ans.append(lit(qx, qy, diag, diag2, row, col))
            lights = turnOff(qx, qy, board, N)
            for l in lights:
                del(board[l])
                x, y = l
                d = x - y
                if d in diag:
                    if l in diag[d]:
                        del(diag[d][l])
                        if len(diag[d]) == 0:
                            del(diag[d])
                d2 = y + x
                if d2 in diag2:
                    if l in diag2[d2]:
                        del(diag2[d2][l])
                        if len(diag2[d2]) == 0:
                            del(diag2[d2])
                if x in row:
                    if l in row[x]:
                        del(row[x][l])
                        if len(row[x]) == 0:
                            del(row[x])
                if y in col:
                    if l in col[y]:
                        del(col[y][l])
                        if len(col[y]) == 0:
                            del(col[y])
        return ans