Leetcode: 407. Trapping Rain Water II

Mistake: Failed to recognize known problem: Shortest Path in a Graph Problem.
Mistake: Did not play with small examples draw invariants.
Can we simulate the problem? We can simulate the water raising from the border to the center of the grid. Every time that water gets to a cell (from the lowest height to the highest height), it will trap water in neighbors that weren't visited and have smaller heights. Time complexity is \(O(n \log n)\) and space is \(O(n)\).

from typing import List


class Solution:
    def trapRainWater(self, h: List[List[int]]) -> int:
        N = len(h)
        M = len(h[0])
        D = [[+1, +0], [-1, +0], [+0, +1], [+0, -1]]

        vis = [[False] * M for _ in range(N)]
        pq = []
        for i in range(N):
            heappush(pq, (h[i][0], i, 0))
            heappush(pq, (h[i][M - 1], i, M - 1))
            vis[i][0] = vis[i][M - 1] = True
        for j in range(M):
            heappush(pq, (h[0][j], 0, j))
            heappush(pq, (h[N - 1][j], N - 1, j))
            vis[0][j] = vis[N - 1][j] = True

        ans = 0
        while pq:
            cur, i, j = heappop(pq)
            for di, dj in D:
                ni = i + di
                nj = j + dj
                if 0 <= ni < N and 0 <= nj < M and not vis[ni][nj]:
                    vis[ni][nj] = True
                    diff = max(cur - h[ni][nj], 0)
                    ans += diff
                    heappush(pq, (max(h[ni][nj], cur), ni, nj))
        return ans


assert (
    Solution().trapRainWater(
        [[1, 4, 3, 1, 3, 2], [3, 2, 1, 3, 2, 4], [2, 3, 3, 2, 3, 1]]
    )
    == 4
)
assert (
    Solution().trapRainWater(
        [
            [3, 3, 3, 3, 3],
            [3, 2, 2, 2, 3],
            [3, 2, 1, 2, 3],
            [3, 2, 2, 2, 3],
            [3, 3, 3, 3, 3],
        ]
    )
    == 10
)