Trapping Rain Water

From LeetCode:

Given an array of non-negative integers that represent an elevation map where the width of each bar is 1, compute the amount of rainwater that can be trapped within it.

Input: [0, 1, 0, 3, 1, 0, 1, 2]
"""
      #
      # ~ ~ ~ #
  # ~ # # ~ # #
0 1 2 3 4 5 6 7
"""
Output: 5

Two Ways of Perceiving Volume

Our goal is to sum up the total amount of water contained in this elevation map. There’s a few ways for us to measure this volume.

An intuitive approach would be to assess how much water each column can support within its unit-width. For every column, we care about the amount of water it can hold directly above itself:

"""
Water: # 1 2 3 # 
       # # 2 3 #
       # # 2 3 #
-----------------
Index: 0 1 2 3 4
"""
# Index 1 holds *1* unit of water
# Index 2 holds *3* units of water
# Index 3 holds *3* units of water

A less straightforward approach would be to divide the volume into rectangles (which includes squares) and associate each rectangle with a column. Summing these blocks of volume would still give us the same answer:

"""
Water: # 1 1 1 # 
       # # 2 2 #
       # # 2 2 #
----------------- 
Index: 0 1 2 3 4
"""
# Index 1 supports *3* units of water
# Index 2 supports *4* units of water

Measuring Volume using Rectangles and Squares

I bring up the second method because, in some ways, this question might remind you of the Largest Rectangle in Histogram problem. After all, in both situations, we’re dealing with an array of non-negative integers that represent unit-width columns of varying height.

Our solution to the Largest Rectangle in Histogram problem required us to maintain a stack of bars of increasing height in order to understand the boundary positions of our candidate rectangles. Can we adapt this strategy to tackle our new problem? Of course!

Instead of preserving a stack of bars of increasing height, we’ll use a stack of bars of decreasing height. When the invariant on the stack is broken (i.e. we encounter a bar that is higher than or equal to the tip of our stack), we’ll pop the stack and compute the area of water associated with the column at that index.

def trap(heights):
    
    active = []
    vol = 0
    
    for i, new_height in enumerate(heights):
        while active and heights[active[-1]] <= new_height:
            popped_idx = active.pop()

            # Ignore cols with no left bound
            if not active: 
                break    

            width = i - active[-1] - 1
            height = (
                min(heights[active[-1]], new_height) -
                heights[popped_idx]
            )
            vol += width * height


        active.append(i)

    return vol

This solution runs pretty well at O(n) space and O(n), but it can be a bit difficult to reason about. Let’s circle back to the first approach we discussed and examine different examples this time.

Measuring Volume by Column

We know that, for any column to hold any water, its left and right sides have to be enclosed by other columns of greater height. Specifically, the enclosing columns will have to be the max-height columns on both sides.

"""
Water:       #  
       # ~ ~ #
       # # ~ # ~ #
-------------------
Index: 0 1 2 3 4 5
"""

• Consider the column at index 2. Column 1 is on next to column 2, but it isn’t the tallest column on the left. The actual tallest column is further down, at index 0.
• On the other side, column 2 is enclosed by column 3.

The other thing we know is that the supported volume will be constrained by the shorter of the two enclosing columns. Since column 0 is of height 2 and column 3 is of height 3, the volume held above column 2 will be min(2, 3) - 0.

To determine supported volumes by brute force, we can do two linear traversals to find the left and right maximum heights:

def trap(heights):

    vol = 0

    for i in range(len(heights)):
        max_left = max(heights[:i] + [0])
        max_right = max(heights[i+1:] + [0])
        vol += max(
            min(max_left, max_right) - heights[i],
            0
        )

    return vol

Running at O(n^2) complexity, this algorithm is far too inefficient. Instead of computing the left and right maximum heights at every iteration, we can reduce the problem space and build on our previous work.

Now imagine that we’ve two pointers, l (left) and r (right), pointing to columns in the elevation map:

"""
Water:          #  
          # ~ ~ #
          # # ~ # ~ #
------------------------
Index:    0 1 2 3 4 5
Pointers:   l     r
Max L Ht: 2 2
Max R Ht:         1 1
"""

• Assume that l is already at index 1, while r is at index 4.
• Suppose we’ve kept track of the maximum heights we’ve seen so far within the left subarray (heights[:l+1]) and the right one (heights[r:]). At this point, the maximum height values are just 2 and 1 respectively.
• The left subarray’s maximum height is currently taller than the right’s, which means that the columns to the right of l could potentially constrain the water held at l. We can’t say anything about the units of water l will support –– we know that it’s at most 1, but we don’t know if it’s 0 or 1.
• On the other hand, the right subarray’s maximum height is currently shorter than the left’s. This means that any other heights to the left of r are irrelevant; the water supported by r will always be limited by it’s right maximum, which we already know! r must support exactly 1 unit of water.

We can think of the elevation map as being composed two subarrays at the edges and expand these subarrays inward by greedily choosing to extend the subarray that has the shorter maximum height.

Of course, we don’t actually need to maintain two arrays in memory, we can just demarcate them with l and r and store their maximums as max_left and max_right:

def trap(heights):

    max_left, max_right = 0, 0
    left, right = 0, len(heights) - 1 
    
    vol = 0
    while left < right:
        # Perpetually update max_left and max_right 
        # as we move "subarrays" inward
        max_left = max(max_left, heights[left])
        max_right = max(max_right, heights[right])
        
        # Left max is shorter or equal to right max,
        # so left max must be the limiting factor 
        if max_left <= max_right:
            vol += max_left - heights[left] 
            left += 1
        
        
        # Right max is shorter than left max,
        # so right max must be the limiting factor
        else: # max_right < max_left
            vol += max_right - heights[right]
            right -= 1
    
    return vol

Given n columns, our algorithm runs in O(1) space and O(n) time.

With that, we’ve solved the Trapping Rain Water problem 🌧📊.

Full Solution

def trap(heights):

    max_left, max_right = 0, 0
    left, right = 0, len(heights) - 1 
    
    vol = 0
    while left < right:
        max_left = max(max_left, heights[left])
        max_right = max(max_right, heights[right])
        
        if max_left <= max_right:
            vol += max_left - heights[left] 
            left += 1
        
        else: # max_right < max_left
            vol += max_right - heights[right]
            right -= 1
            
    
    return vol