Grid Illumination Python Solution with Tests

lcpy gen -n 1001   # by problem number
lcpy gen -s grid_illumination   # by problem name

Input: n = 5, lamps = [[0,0],[4,4]], queries = [[1,1],[1,0]]
Output: [1,0]
Explanation: We have the initial grid with all lamps turned off. In the above picture we see the grid after turning on the lamp at grid[0][0] then turning on the lamp at grid[4][4].
The 0th query asks if the lamp at grid[1][1] is illuminated or not (the blue square). It is illuminated, so set ans[0] = 1. Then, we turn off all lamps in the red square.

Input: n = 5, lamps = [[0,0],[4,4]], queries = [[1,1],[1,1]]
Output: [1,1]

Input: n = 5, lamps = [[0,0],[0,4]], queries = [[0,4],[0,1],[1,4]]
Output: [1,1,0]

class Solution:
    # Time: O(L + Q) where L = len(lamps), Q = len(queries) (9 turns per query)
    # Space: O(L)
    def grid_illumination(
        self, n: int, lamps: list[list[int]], queries: list[list[int]]
    ) -> list[int]:
        lit: set[tuple[int, int]] = set()
        for r, c in lamps:
            lit.add((r, c))
        rows: dict[int, int] = {}
        cols: dict[int, int] = {}
        diag: dict[int, int] = {}
        anti: dict[int, int] = {}
        for r, c in lit:
            rows[r] = rows.get(r, 0) + 1
            cols[c] = cols.get(c, 0) + 1
            diag[r - c] = diag.get(r - c, 0) + 1
            anti[r + c] = anti.get(r + c, 0) + 1

        result: list[int] = []
        for r, c in queries:
            is_lit = (
                rows.get(r, 0) > 0
                or cols.get(c, 0) > 0
                or diag.get(r - c, 0) > 0
                or anti.get(r + c, 0) > 0
            )
            result.append(1 if is_lit else 0)
            for dr in (-1, 0, 1):
                for dc in (-1, 0, 1):
                    lamp = (r + dr, c + dc)
                    if lamp in lit:
                        lit.remove(lamp)
                        rows[lamp[0]] -= 1
                        cols[lamp[1]] -= 1
                        diag[lamp[0] - lamp[1]] -= 1
                        anti[lamp[0] + lamp[1]] -= 1
        return result