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

Step-by-step optimisation.

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

Brute force

Determine the water on i-th position.

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.

Use two pointers to scan the input.
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])
                ret += max_r-height[r]
                r -= 1
                max_r = max(max_r,height[r])
        return ret


Notify of
Inline Feedbacks
View all comments