Problem statement

https://leetcode.com/problems/trapping-rain-water-ii/

Solution

Quite difficult problem, extension of Problem 0042 Trapping Rain Water (which is already marked as hard). The clue point here is to understand, what does trapping water means. First of all, we have boundary, which can not be trapped. Also, for each height we can evaluate in O(mn) operations 0-1 matrix with indicator more or less and equal that this height and then use BFS to found regions of 0, fully surrounded by $1$. In this way we can have O(mnH) complexity.

There is smarter algorithm, using heap. We discussed that border can keep any water, so let us put it into heap. Now, we are interested in the smallest element, so:

  1. We pop it out of our heap, so we have height, x, y the smallest height so far in our heap, and coordinates.
  2. For all neighbors dx, dy of this points, which are not visited yet (we keep set of visited cells), if height of this neighbor is more than current height, it means we can not trap any water. If it is less, we can, and we add height - A[dy][dx] amount of water to our final answer.
  3. Finally, what we push into our heap is max(A[dy][dx], height) and mark (dx, dy) as visited.

Complexity

Time complexity of this approach is O(mn * log(mn)), because we traverse each element only two times: when we put it into heap and when we remove it.

Code

class Solution:
    def trapRainWater(self, A):
        m, n, ans = len(A), len(A[0]), 0
        V = set(product([0, n-1], range(m))) | set(product(range(n), [0, m-1]))
        h = [(A[y][x], x, y) for x, y in V]
        heapify(h)
        
        while h:
            height, x, y = heappop(h)
            for dx, dy in (x+1, y), (x, y+1), (x-1, y), (x, y-1):
                if 0 <= dx < n and 0 <= dy < m and (dx, dy) not in V:
                    ans += max(0, height - A[dy][dx])
                    heappush(h, (max(A[dy][dx], height), dx, dy))
                    V.add((dx, dy))
        
        return ans

Remark

There is really good youtube visualization of this algorithm https://www.youtube.com/watch?v=cJayBq38VYw