LeetCode 42 — Trapping Rain Water

The problem can be found here.

A very bad habit of mine, when encountered with an array problem such as this, is that my mind immediately starts looking for one-pass solutions without conceptualizing even a brute-force solution first, knowing that a brute-force in this case almost always guarantees an O(n²) time complexity. While this can prove fatal sometimes, fortunately, a brute-force algorithm is accepted for Python on LeetCode for this problem; but that is beside the point. A one-pass solution might be intuitive as one becomes more familiar with array problems in general, perhaps it would serve someone like myself better to think of a brute-force algorithm first before trying to optimize it, especially during interviews where the ability to explain your thought process is of vital importance.

Brute Force

For a newbie like me, a brute-force solution should be the most intuitive one where you simply do as you are told. Let’s look at the provided example, which can also be found below:

The above map is represented by the array [0,1,0,2,1,0,1,3,2,1,2,1]; black represents terrain elevation at each index, and blue represents the rainwater trapped

How do we find the amount of water trapped at index 2? We simply take the two “walls” on both sides of the index, and the lower height of both walls determines how much water can be retained, which in this case is 1, being the lower of 1 (at index 1) and 2 (at index 3).

But what happens when we’re dealing with a “pool” of water (such as at index 4, 5, 6) instead of just a single column at index 2? Does the pattern still hold? Let’s look at index 4: the “walls” on either side of it are at index 3 and index 7 — notice how in order to qualify as a “wall”, the elevation at that index must be higher than the current index, and index 3 and 7 are the closest walls to index 4. Among these two walls, the lower of the two are at index 3 of heights 2; therefore the water retained at index 4 is also at level 2, although we need to subtract its base elevation of 1 to get the actual water amount 1. This pattern holds for water mounts at index 5, 6, and 9, respectively.

We can then generalize our brute-force algorithm as follows:

At each index, the amount of water trapped is equal to the lower heights of two walls on its either side, minus the base elevation at that index.

More succinctly, let’s write it this way:

water[current] = min(left_wall_height, right_wall_height) —    
                 height[current]

To get the left_wall_height and the right_wall_height, we only need to perform a maximum operation on all the elements to the left of the current index, and all the elements to the right of the current index, respectively. In other words, they can be represented as:

left_wall_height[i] = max(height[:i])
right_wall_height[i] = max(height[i+1:])

Now we’re ready to write this down in code. But before we do that, some questions to think about:

1. What do we do with edges (i.e. index 0, index len(heights)-1?)
2. Are there any edge cases or bad inputs we need to be wary of?

res = 0 
for i in range(1, len(height)-1):
    h = height[i]
    left = max(height[:i])
    right = max(height[i+1:])
    if min(left, right) > h:
        res += min(left, right) - h
return res

Brute Force Optimization

It soon becomes apparent that we are doing the max operation on either side of any given index exactly n times, which is wasteful in terms of time complexity; we should already know everything we need to know by iterating through the input once. Is there a way we can trade space for speed?

Let’s rethink our solution above. What we need to know to arrive at the water amount at each index, is the height of the walls to its left and right, as well as its own height. Is there a way we can record the left walls and the right walls at each index in some kind of data structure, and then simply access that information by the index we are trying to calculate water amount at?

Dynamic Programming

DP solutions are a little trickier to conceptualize right off the bat, but since we already know from the brute force solution above what information we need to save, it should become apparent that, as part of the DP solution, we will need to build an array that stores all tallest left walls at any given index, as well as an array that stores all tallest right walls at any given index.

What do I mean by this? Let’s take another look at the example we’re already familiar with.

[0,1,0,2,1,0,1,3,2,1,2,1]

At index 0, because there is nothing to its left, let’s say the tallest wall to its left has a height of 0.

At index 1, the tallest wall to its left has a height of 0 (at index 0).

At index 2, the tallest wall to its left has a height of 1 (at index 1).

At index 3, the tallest wall to its left has a height of 1 (at index 1).

At index 4, the tallest wall to its left has a height of 2 (at index 3).

And so on.

Thus, we can record the tallest left wall at each index in an array such as the one shown below:

We do the same for the walls to the right, except this time we need to look at each index from the right to the left.

At index 11, because there is nothing to its right, let’s say the tallest wall to its right has a height of 0.

At index 10, the tallest wall to its right has a height of 1 (at index 11).

At index 9, the tallest wall to its right has a height of 2 ( at index 10).

And so on.

Similarly, heights of the tallest right walls are stored in the array below:

After this, we should be able to do what we did in the brute-force solution without having to search the whole input array for the tallest left/right wall every time we move to a new index.

The code below should be fairly intuitive at this point:

res = 0
max_left, max_right = [0 for _ in height], [0 for _ in height]
for i in range(1, len(height)): # We start at index 1
    # The maximum height at each index is the height of the index 
    # to its immediate left, or the tallest wall at that index    
    max_left[i] = max(height[i-1], max_left[i-1])
for i in range(len(height)-2, -1, -1): # We start at second    
                                       # rightmost index and go 
                                       # backwards
    max_right[i] = max(height[i+1], max_right[i+1])
# The below should be almost identical to the brute force solution
for i in range(1, len(height)-1):
    h = height[i]
    left = max_left[i]
    right = max_right[i]
    if min(left, right) > h:
        res += min(left, right) - h
return res

Dynamic Programming Optimization

It is not hard to notice how we only accessed the information regarding tallest walls exactly once at each index, therefore this O(n) space solution can be further optimized down to an O(1) space solution. But how exactly do we do that?

Notice how we stored all tallest left walls as we iterated the input array from left to right, and stored all tallest right walls as we iterated from right to left. If only we can get the tallest left and right walls at each index and calculate the amount of water at that index at the same time, then we won’t need to store all that information ahead of time and therefore won’t be needing O(n) space.

So how do we get the tallest left wall and the tallest right wall, while the former requires us to iterate from left to right, and the latter right to left?

The answer should not be too hard to come by.

Two Pointers

The two pointers technique is something we should always keep in our toolbox when dealing with array problems like this. It will come in handy especially when we need to iterate an array in two directions or search for pairs in an array.

The only thing that’s going to change from our two solutions above, other than how we look for tallest walls, is that we can no longer get the water amount at each index in order, from left to right. This is the primary reason why I first found this solution to be less intuitive than the other ones; perhaps like everything else it just needs some getting used to.

Why can’t we get the water amount simply from left to right? Let’s look at the example again.

[0,1,0,2,1,0,1,3,2,1,2,1]

We already know how to get the heights of the tallest left/right walls; we iterate from the left and update the tallest left wall as we move our left pointer, and we iterate from the right and update the tallest right wall as we move our right pointer. So the left pointer must start from index 1, and the right pointer from index len(heights)-2 (second rightmost index).

But is such a wide bound going to help us narrow down the water amount at index 1, which is the index we started with from the previous two solutions? Is it going to help us with index 2?

If we observe closely as the two pointers move toward the center, we might see a pattern.

If we start our left pointer from index 1, then by the time we get to index 2 as shown above, we will have already found that the max_left for index 2 is at index 1. The same goes for the right pointer; when we start at the second rightmost index (right), we know that the index immediately to its right has the highest wall (max_right).

How does knowing this help us understand the water levels at both left and right? We know that the water level at either index is bound by the lower heights of the two walls (max_left, max_right). It shouldn’t matter that these two walls are far apart from each other; we can easily imagine a scenario where there is no elevation between max_left and max_right, and the rule still holds.

How do we move our pointers, then? From the above, we know that the water level is dictated by the lower of max_left and max_right. Therefore, once we observe that max_left < max_right, we fill the left index up to max_left, and advance the left pointer; and vice versa. As for when they are the same, either works.

This can then be summarized as below:

max_left = max(max_left, height[left-1]
max_right = max(max_right, height[right+1]
if max_left <= max_right:
    if max_left > height[left]:
        res += max_left - height[left]
    left += 1
else:
    if max_right > height[right]:
        res += max_right - height[right]
    right -= 1

To further illustrate: from the last state, we have our max_left = max_right. According to the code above, we move our left pointer, all the way until we have a max_left that is higher than max_right. We first find this left pointer at index 4:

Now max_right becomes the lower wall, so we know that at the right position, the max water level is max_right, which is 1. We can see this is true if we imagined a scenario where the elevation at right is 0, instead of 2 like the example above.

Now we just need to initialize the pointers, initial max_left, max_right, and put the above inside a while loop, and we have our code:

res = 0
max_left, max_right = 0, 0
left, right = 1, len(height)-2
while left <= right:
    max_left = max(max_left, height[left-1])
    max_right = max(max_right, height[right+1])
    if max_left <= max_right:
        if max_left > height[left]:
            res += max_left - height[left]
        left += 1
    else:
        if max_right > height[right]:
            res += max_right - height[right]
        right -= 1     
return res

Stack

The motivation for a stack-based solution should be fairly intuitive. We know that water can potentially accumulate wherever the elevation increases, therefore if we keep a stack that adds new indexes while comparing its height against the height of the index at the stack top, we can easily calculate the amount of water when the height at the current index is higher than at the stack top.

But what exactly do we do when we encounter an index higher than stack top? Because of the nature of water accumulation — the fact that we need both a left wall and a right wall that are at least one index apart from each other for there to be water in the middle — we know that the current index would be our right wall, the stack top is where the water potentially accumulates, then after we pop the stack top, the very next one in the stack would be our candidate for the left wall.

Because we are no longer adding water amount column by column, it is necessary to know not only the height but also the width of any water accumulation, therefore we add indexes instead of heights into the stack.

The code is as follows:

res = 0
stack = []
for i in range(len(height)):
    while stack and height[i] > height[stack[-1]]:
        h = height[stack.pop()]
        if not stack:
           break
        res += (min(height[i], height[stack[-1]]) - h) * (i - stack[-1] - 1)
    stack.append(i)
return res