[Leetcode]42. Trapping Rain Water

Given n non-negative integers representing an elevation map where the width of each bar is 1, compute how much water it is able to trap after raining.
Example:
Input: height=[0,1,0,2,1,0,1,3,2,1,2,1] Output: 6

This is a great Leetcode question and solving it reveals a clear path from the brute force approach to optimised ones using dynamic programming and two pointers.

Note that this post gets inspiration from @HuaHua, and his video is attached at the end.

Brute force

As shown above, we can observe that, for each position, the water saved is determined by the shorter one of max_l and max_r, where max_l and max_r are the highest bars in its lefthand and righthand side. Following this intuition, we can implement the solution with O(n^2) time complexity: for each i, it saves water: min(max(height[0...i]),max(height[i...n-1]))-height[i].

class Solution:
    def trap(self, height: List[int]) -> int:
        ret = 0
        for i in range(1,len(height)-1):
            ret += max(min(max(height[:i]),max(height[i+1:]))-height[i],0)
        return ret

Dynamic programming

As it’s obvious that we keep looking for the highest bars in left- and right-hand side, we would be able to improve its time efficiency if we store the result along the way: instead of calling max(height[0..i]) for each i we could simply call max(height[i],max_l[i-1]), where max_l[i-1] records the highest bars in height[0,i-1].

class Solution:
    def trap(self, height: List[int]) -> int:
        if not height:
            return 0
        # prepare max_l and max_r
        l,r = [-1] * len(height),[-1] * len(height)
        l[0] = height[0]
        r[-1] = height[-1]
        ret = 0
        for i in range(1,len(height)):
            l[i] = max(l[i-1],height[i])
        for i in range(len(height)-2,-1,-1):
            r[i] = max(r[i+1],height[i])

        for i in range(1,len(height)-1):
            ret += min(l[i],r[i])-height[i]
        return ret

Two pointers

Finally, if you notice that the water saved on i-th bar is determined by the shorter one of max_l and max_r, we can replace the arrays with two pointers. For example, in the case shown below, we know max_l < max_r, then water saved on l is already determined by max_l; or if max_l > max_r, r is determined.

class Solution:
    def trap(self, height: List[int]) -> int:
        if len(height) < 2:
            return 0
        l,max_l = 0,height[0]
        r,max_r = len(height)-1,height[-1]
        ret = 0
        while l < r:
            if max_l < max_r:
                ret += max_l-height[l]
                l += 1
                max_l = max(max_l,height[l])
            else:
                ret += max_r-height[r]
                r -= 1
                max_r = max(max_r,height[r])
        return ret

Reference