Python solution using heap + dfs beats 99%


  • 2
    T

    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[0]) < 3: return 0
            m, n, bd = len(heightMap), len(heightMap[0]), []
            for x, y in zip([0]*n+[m-1]*n+range(1,m-1)*2, range(n)*2+[0]*(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
    

  • 0
    V

    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
    

  • 0
    T

    @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.


  • 0
    V

    @TeXnician Yeah.. Looks like there is a discrepancy between problem description and the test cases.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.