Leetcode 407 Solution

class Solution:
    def trapRainWater(self, heights: List[List[int]]) -> int:
        # write your code here
        n = len(heights)
        if n <= 1:
            return 0

        m = len(heights[0])
        if m <= 1:
            return 0

        heap = []
        seen = set()
        for i in range(n):
            heapq.heappush(heap, (heights[i][0], i, 0))
            heapq.heappush(heap, (heights[i][m - 1], i, m - 1))
            seen.add((i, 0))
            seen.add((i, m - 1))
        for j in range(m):
            heapq.heappush(heap, (heights[0][j], 0, j))
            heapq.heappush(heap, (heights[n - 1][j], n - 1, j))
            seen.add((0, j))
            seen.add((n - 1, j))

        level = 0

        ans = 0
        while heap:
            level, i, j = heapq.heappop(heap)
            ans += max(0, level - heights[i][j])

            neis = []
            if i > 0 and (i - 1, j) not in seen:
                neis.append((i - 1, j))
            if i < n - 1 and (i + 1, j) not in seen:
                neis.append((i + 1, j))
            if j > 0 and (i, j - 1) not in seen:
                neis.append((i, j - 1))
            if j < m - 1 and (i, j + 1) not in seen:
                neis.append((i, j + 1))

            for next_i, next_j in neis:
                seen.add((next_i, next_j))
                heapq.heappush(heap, (max(level, heights[next_i][next_j]), next_i, next_j))

        return ans