题目链接:1001. 网格照明
使用四个列表分别记录当前行、列、对角线、反对角线上灯的数目。
遍历一次灯的列表,初始化这四个列表和灯的集合;再遍历一次查询的列表,每次查询,将这个点及周围共 9 个点位置的灯删除,并修改相应的列表。
from collections import defaultdict
from typing import List
class Solution:
def gridIllumination(self, n: int, lamps: List[List[int]], queries: List[List[int]]) -> List[int]:
lamps = set(map(tuple, lamps))
rows, cols, diags, neg_diags = defaultdict(int), defaultdict(int), defaultdict(int), defaultdict(int)
for x, y in lamps:
rows[x] += 1
cols[y] += 1
diags[x - y] += 1
neg_diags[x + y] += 1
ret = []
for x, y in queries:
ret.append(1 if rows[x] or cols[y] or diags[x - y] or neg_diags[x + y] else 0)
for i, j in [
(x - 1, y - 1), (x - 1, y), (x - 1, y + 1),
(x, y - 1), (x, y), (x, y + 1),
(x + 1, y - 1), (x + 1, y), (x + 1, y + 1),
]:
if (i, j) not in lamps:
continue
lamps.remove((i, j))
rows[i] -= 1
cols[j] -= 1
diags[i - j] -= 1
neg_diags[i + j] -= 1
return ret