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 rain. For example,

Input: [0,1,0,2,1,0,1,3,2,1,2,1] 
Output: 6

trapping rain water

Basic thought

We know water stays at the highest level it is able to, and it always maintains the same flat surface. Using this, we can infer that we need to find holes in the elevation where water would be able to rest at a level. To calculate how much water these holes would need to store, we can see that we need to have elevations on both sides, and we also need to track how much space a particular hole would be able to trap the water. Fig below is an example of a hole which holds 4 units of water.

trap water

Upon further breakdown of these holes, we can notice that we do not need to track the entire hole to find the capacity it holds, but we can parse each unit of the hole individually. Thus, the amount of water in each unit of the hole is

min(leftHeight, rightHeight) - currentUnitHeight.

What remains now is to calculate the leftHeight and the rightHeight. We could parse through them individually to find these out, but we can see a general pattern here: The highest elevation to the left inclusive of the current unit will become the leftHeight, and the highest elevation to the right inclusive of the current unit will become the rightHeight. The problem has been greatly simplified into maintaining track of highest heights on both sides of every unit.

Brute force solution

From our observations in the previous section, the simplest brute force approach is to calculate the highest elevation on both sides of every unit, and sum them up together. Each unit takes O(n) time with this approach, and there are n units to calculate in total. Thus, this approach will take O(n2) time.

Show me the brute force implementation

class Solution:
    def trap(self, height: List[int]) -> int:
        if len(height) < 3:
            return 0
        ans = 0
        for i in range(len(height)):
            left_max = 0
            right_max = 0
            for j in range(i + 1):
                left_max = max(left_max, height[j])
            for j in range(i, len(height)):
                right_max = max(right_max, height[j])
            ans += min(left_max, right_max) - height[i]
        return ans

Dynamic Programming approach

As we can see from the brute force solution, we calculate the leftHeight and the rightHeight multiple times for the same node, i.e. the problem has overlapping subproblems. Thus a dynamic programming approach should optimize the brute force approach further. We can store the leftHeight and rightHeight elements till each index we have iterated, and thus the water storage calculation for each unit will now take O(1) time. Overall, this approach passes over the array thrice, and thus has a runtime of O(n) with a space complexity of O(n).

Show me the dynamic programming implementation

class Solution:
    def trap(self, height: List[int]) -> int:
        if len(height) < 3:
            return 0
        ans = 0
        left_max = [0] * len(height)
        right_max = [0] * len(height)
        left_max[0] = height[0]
        right_max[-1] = height[-1]
        for i in range(1, len(height)):
            left_max[i] = max(height[i], left_max[i - 1])
        for i in reversed(range(len(height) - 1)):
            right_max[i] = max(height[i], right_max[i + 1])
        for i in range(1, len(height) - 1):
            ans += min(left_max[i], right_max[i]) - height[i]
        return ans

Stacks

Since we need to keep track of the highest elevations up to a point, stacks are a good approach to perform this operation in one pass of the array. The basic idea is since we need to store the largest elevations in the stack, as we iterate through the array, we can find the amount of water stored till the currentHeight is higher than elements of the stack and move on to the next value, i.e. water stored will always be of a hole shape, thus we can find the amount of water that can be stored between two high values. This operation passes over the array once, and each element can only have two operations maximum: Pushing and popping from the stack, thus its time complexity is O(n). The space complexity will be O(n), in case the entire array is stored on the stack.

Show me the stack implementation

class Solution:
    def trap(self, height: List[int]) -> int:
        if len(height) < 3: return 0 ans = 0 stack = [] for i in range(len(height)): while len(stack) > 0 and height[i] > height[stack[-1]]:
                top = stack.pop()
                if len(stack) == 0:
                    break
                # Distance between the larger value still in the 
                #stack with a hole the height of the top element 
                #and the current element
                distance = i - stack[-1] - 1
                # Water that can be stored is smaller 
                # heights between these bounds, and the height 
                # of the intermediate region between top of the stack 
                # and current index
                curr_height = min(height[i], height[stack[-1]]) - height[top]
                ans += distance * curr_height
            stack.append(i)
        return ans

Simple Optimization

Another way to approach the problem is that since we need to find the maximum elevations on either side to calculate the current water stored, we can calculate the global maximum in one pass, and once we have the index for the same, we can iterate to it and from it to the rest of the array knowing that we have one bounded measurement which is the highest elevation in the array. The time complexity for this operation is O(n), but there are two passes over the array(once to calculate global maximum, and once to calculate the water amount). The space complexity is O(1).

Show me the implementation

class Solution:
    def trap(self, height: List[int]) -> int:
        if len(height) < 3:
            return 0
        ans = 0
        gMax = 0 # Global Max
        for i in range(len(height)):
            gMax = i if height[gMax] < height[i] else gMax
        lMax = 0 # Left max yet
        for i in range(1, gMax):
            lMax = i if height[lMax] <= height[i] else lMax
            ans += max(0, height[lMax] - height[i])
        lMax = len(height) - 1 # Right max yet
        for i in reversed(range(gMax, len(height) - 1)):
            lMax = i if height[lMax] <= height[i] else lMax
            ans += max(0, height[lMax] - height[i])
        return ans

This article is contributed by Khushman Patel