bd stores the current boundary(boundaries) of the water container.
Note that since the heights are always positive, we can set
heightMap[x][y] = 0 as
visited. No need to use another O(mn) space.
from heapq import * class Solution(object): def trapRainWater(self, heightMap): """ :type heightMap: List[List[int]] :rtype: int """ if len(heightMap) < 3 or len(heightMap) < 3: return 0 m, n, bd = len(heightMap), len(heightMap),  for x, y in zip(*n+[m-1]*n+range(1,m-1)*2, range(n)*2+*(m-2)+[n-1]*(m-2)): heappush(bd, [heightMap[x][y], x, y]) heightMap[x][y] = -1 res = 0 while bd: h, x, y = heappop(bd) res = self.dfs(h, x, y, bd, heightMap, m, n, res) return res def dfs(self, h, x, y, bd, heightMap, m, n, res): for dx, dy in zip([-1,1,0,0], [0,0,-1,1]): if x+dx > 0 and x+dx < m - 1 and y+dy > 0 and y+dy < n - 1 and heightMap[x+dx][y+dy] >= 0: tmpH = heightMap[x+dx][y+dy] heightMap[x+dx][y+dy] = -1 if tmpH <= h: res += h - tmpH res = self.dfs(h, x+dx, y+dy, bd, heightMap, m, n, res) else: heappush(bd, [tmpH, x+dx, y+dy]) return res
It gives wrong answer in the following case:
Input: [[9,9,9,9,9,9,8,9,9,9,9],[9,0,0,0,0,0,1,0,0,0,9],[9,0,0,0,0,0,0,0,0,0,9],[9,0,0,0,0,0,0,0,0,0,9],[9,9,9,9,9,9,9,9,9,9,9]] Output: 7 Expected: 215
@voiceup The test cases are changed. As stated in the problem description, all the elements in the matrix should be positive. But this input has zero elements. I just modified it by setting the height to be -1 and it passes the new cases.
@TeXnician Yeah.. Looks like there is a discrepancy between problem description and the test cases.
Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.