Problem statement

https://binarysearch.com/problems/Rain-Catcher-Sequel/

Solution

Equal to Leetcode 407. Trapping Rain Water II

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 solve(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