Leetcode: 42. Trapping Rain Water

Problem Statement

How can we extend the solution for \(i\) to \(i+1\)? There is nothing to do if we have just one or two walls. With three walls \(h_1, h_2\) and \(h_3\), there is no water trapped if \(h_1h_2>h_3\), but there is water trapped if \(h_1>h_2\) and \(h_3>h_2\). In the last case, we can compute the amount of water trapped \((\min(h_3,h_1) - h_2) \times (i_3 - i_1 - 1)\). After that, we can remove \(a_2\), add \(a_3\) and continue the process with the next height. We can use Monotonic Stack to keep track of the heights and solve the problem in time and space \(O(n)\).

from typing import List


class Solution:
    def trap(self, height: List[int]) -> int:
        ans = 0
        st = []
        for i in range(len(height)):
            while st and height[st[-1]] <= height[i]:
                c = st.pop()
                if st:
                    h = min(height[i], height[st[-1]]) - height[c]
                    w = i - st[-1] - 1
                    ans += h * w
            st.append(i)
        return ans


assert Solution().trap([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]) == 6
assert Solution().trap([4, 2, 0, 3, 2, 5]) == 9