Maximum Number of Points From Grid Queries

lcpy gen -n 2503   # by problem number
lcpy gen -s maximum_number_of_points_from_grid_queries   # by problem name

Input: grid = [[1,2,3],[2,5,7],[3,5,1]]
queries = [5,6,2]
Output: [5,8,1]
Explanation: The diagrams above show which cells we visit to get points for each query.

Input: grid = [[5,2,1],[1,1,2]]
queries = [3]
Output: [0]
Explanation: We can not get any points because the value of the top left cell is already greater than or equal to 3.

import heapq


class Solution:
    # Time: O(m*n*log(m*n) + k*log k)
    # Space: O(m*n)
    def max_points(self, grid: list[list[int]], queries: list[int]) -> list[int]:
        rows, cols = len(grid), len(grid[0])
        visited = [[False] * cols for _ in range(rows)]
        visited[0][0] = True
        heap: list[tuple[int, int, int]] = [(grid[0][0], 0, 0)]
        count = 0
        counts: dict[int, int] = {}
        for query in sorted(set(queries)):
            while heap and heap[0][0] < query:
                _, i, j = heapq.heappop(heap)
                count += 1
                for ni, nj in ((i + 1, j), (i - 1, j), (i, j + 1), (i, j - 1)):
                    if 0 <= ni < rows and 0 <= nj < cols and not visited[ni][nj]:
                        visited[ni][nj] = True
                        heapq.heappush(heap, (grid[ni][nj], ni, nj))
            counts[query] = count
        return [counts[query] for query in queries]