Trapping rain water with a Python list as a stack

Why this matters

Some questions describe a landscape with a row of bars, and ask how much rain would stay trapped between them. The bars are the ground, the gaps between tall bars are the low places, and the water sits in those low places until it spills over the lower of the two walls around it. The question is common in interviews because the easy first answer looks at each bar on its own, and the faster answers take some care to get right.

This lesson builds the stack method from the idea of filling water in layers. It then compares the stack method with a two-pointer scan, which reaches the same answer by a different route. The stack method uses the same kind of monotonic stack as the two previous lessons in this module, where the stack holds items in one order. Here the stack holds bars that are still waiting for a taller bar on their right.

Words you will need

• Bar: one column of the landscape. Its height is a whole number, and its position is its index, counted from 0.
• Wall: a bar that stops the water from running away. The water in a gap is held by a wall on the left and a wall on the right.
• Basin: a low gap between two walls. The bottom of the basin is the lowest bar inside the gap.
• Level: the height of the lower wall around a basin. Water fills up to that height, and no higher.
• Depth: how far the water rises above the bottom of the basin. It is the level minus the height of the bottom bar.
• Width: how many bars wide a basin is. It counts the bars between the two walls, not the walls themselves.
• Two pointers: two index variables, one at each end of the list, that move inward toward each other.
• Brute force: the simplest correct method, which checks every bar against every other bar. It is slow, but it is a good reference for checking a faster method.
• List as a stack: a Python list where append adds to the end and pop() removes from the end. The end of the list is the top of the stack. Python's standard library has no plain stack class, so a list is the usual choice.

Start with the picture

Take the bars [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]. Draw them as columns, with the tallest one reaching height 3. Rain falls on the whole row and then stays put. Each gap that has walls on both sides collects water. A gap with no wall on one side lets the water run off, so it holds nothing.

The first gap is at index 2, where the height is 0. Its left wall is index 1, with height 1, and its right wall is index 3, with height 2. The water in that gap rises to height 1, which is the lower of the two walls. The gap is one bar wide, and the water is one unit deep above the bottom at height 0. So that gap holds 1 unit of water.

Each gap is a rectangle of water. The water in a basin has a width and a depth, and the amount of water is the width times the depth. If you add up the rectangles of all basins, you have the total.

Filling the water level by level

Another way to see the same answer is to fill the landscape with water one horizontal layer at a time. Start with the layer at height 1. Which bars are at height 0? Those are the only places where water can sit at height 1, and only if a wall on each side is at least height 1. Count the bars at height 0 that are between two walls of height 1 or taller. Each of those counts as one unit of water in this layer. Then move to height 2, and count the bars below 2 that are between two walls of height 2 or taller. Repeat for every height up to the tallest bar.

This layer view is correct, but counting layers one by one is slow when the bars are tall. The stack method does the same counting in a different order. It works basin by basin, and each basin is filled to its full level in one step. The next sections show how.

The brute force first

For each bar, find the tallest bar on its left, including itself, and the tallest bar on its right, including itself. The water above that bar is the smaller of the two maximums, minus the bar's own height. Add up the amounts for all bars.

def trap_brute_force(h):
    water = 0
    for i in range(len(h)):
        left_max = max(h[:i + 1])
        right_max = max(h[i:])
        water += min(left_max, right_max) - h[i]
    return water

h[:i + 1] is a slice. It makes a new list with the items from index 0 up to, but not including, index i + 1. So it holds the bars from the left end through bar i. h[i:] holds the bars from i to the end. The built-in max and min functions return the largest and smallest item in a list or in a group of values.

Why is the smaller of the two maximums the right level? Water above a bar rises until it reaches the lower of the two walls that hold it. The tallest bar on the left is the highest wall the water could meet from that side. The tallest bar on the right is the same from the other side. The water stops at the smaller of the two, because it would spill over the lower one first.

The brute force makes a slice and a max call for every bar. For n bars that is about n times n work, which is O(n squared). It is correct, and it is the reference answer for the faster methods.

The stack method

A stack can keep the bars that are still waiting for a wall on their right. Scan from left to right. Keep the stack in order: from the bottom of the list to the top, the bars get shorter. When a new bar is taller than the bar on top of the stack, the top bar is the bottom of a basin. Its right wall is the new bar. Its left wall is the bar left under it in the stack, because that bar is the nearest taller bar on the left.

Pop the top bar and compute the water above it. The level is the lower of the two walls. The width is the number of bars between the two walls. The depth is the level minus the height of the popped bottom. Add width times depth to the total. Then check the next top, because the new bar may be taller than that one too. Repeat until the stack is empty or the new bar is not taller than the top. Then push the new bar.

def trap_with_stack(h):
    stack = []
    water = 0
    for i, height in enumerate(h):
        while stack and height > h[stack[-1]]:
            bottom = stack.pop()
            if not stack:
                break
            left = stack[-1]
            width = i - left - 1
            level = min(height, h[left]) - h[bottom]
            water += width * level
        stack.append(i)
    return water

The stack holds indices, not heights, as in the daily temperatures lesson. The width needs index subtraction, and the heights are read back from the index with h[stack[-1]]. stack[-1] means the last item in the list. Python counts negative indexes from the end, so -1 is the top of the stack.

enumerate(h) gives a pair on each step: the index i and the height height. A plain for x in h loop would give only the height, and then the code would not know the index.

Python lists have no peek method. Reading stack[-1] does the same job as peek: it looks at the top without removing it.

The if not stack: break lines matter. If the popped bar has nothing on its left, there is no left wall, so the water runs away and the bar holds nothing. The loop stops, and the new bar is pushed. Bar 3 in the first example is this case: no taller bar sits to its left, so when bar 7 pops it, nothing holds the water.

The next program runs all three methods on the example and on three small cases. It has the data built in, so you can run it without typing input.

The output is:

bars     [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
brute    6
stack    6
two      6
plateau  10
valley   9
staircase 0

Worked example 1: the landscape from the picture

The trace below is a version of the stack function that prints one row for each bar. In each row, before is the stack before the bar is handled. A Python list prints with the bottom first, so the last number in each before list is the top. The notes part names each basin that was filled and how much water it added. level in a note is the height of the lower wall.

def trace_stack(h):
    stack = []
    water = 0
    for i, height in enumerate(h):
        notes = []
        before = list(stack)
        while stack and height > h[stack[-1]]:
            bottom = stack.pop()
            if not stack:
                notes.append(f"bar {bottom} has no left wall")
                break
            left = stack[-1]
            width = i - left - 1
            wall_level = min(height, h[left])
            level = wall_level - h[bottom]
            water += width * level
            notes.append(f"basin at {bottom} walls {left},{i} width {width} level {wall_level} adds {width * level}")
        stack.append(i)
        note_text = "; ".join(notes) if notes else "-"
        print(f"i {i:2} h {height} | before {str(before):<18} | {note_text} | total {water}")
    return water

The output of trace_stack on the example is:

i  0 h 0 | before []                 | - | total 0
i  1 h 1 | before [0]                | bar 0 has no left wall | total 0
i  2 h 0 | before [1]                | - | total 0
i  3 h 2 | before [1, 2]             | basin at 2 walls 1,3 width 1 level 1 adds 1; bar 1 has no left wall | total 1
i  4 h 1 | before [3]                | - | total 1
i  5 h 0 | before [3, 4]             | - | total 1
i  6 h 1 | before [3, 4, 5]          | basin at 5 walls 4,6 width 1 level 1 adds 1 | total 2
i  7 h 3 | before [3, 4, 6]          | basin at 6 walls 4,7 width 2 level 1 adds 0; basin at 4 walls 3,7 width 3 level 2 adds 3; bar 3 has no left wall | total 5
i  8 h 2 | before [7]                | - | total 5
i  9 h 1 | before [7, 8]             | - | total 5
i 10 h 2 | before [7, 8, 9]          | basin at 9 walls 8,10 width 1 level 2 adds 1 | total 6
i 11 h 1 | before [7, 8, 10]         | - | total 6

Read the row for i = 7, the bar of height 3. The stack before it is [3, 4, 6], so the top is bar 6, of height 1. Bar 7 is taller, so bar 6 is a basin bottom. Its left wall is bar 4, of height 1, and its right wall is bar 7. The gap between the walls holds bars 5 and 6, so the width is 7 minus 4 minus 1, which is 2. The level is the lower wall, min(3, 1), which is 1. The depth is the level minus the bottom's height, 1 minus 1, which is 0. So this basin adds 0. Bar 6 is popped.

Next the top is bar 4, of height 1. Bar 7 is still taller. The walls are now bar 3 on the left (height 2) and bar 7 on the right (height 3). The gap holds bars 4, 5 and 6, so the width is 3. The level is min(3, 2), which is 2. The depth is 2 minus the bottom's height of 1, which is 1, so the basin adds 3 * 1 = 3 units. Bar 4 is popped.

The top is now bar 3, of height 2, which is shorter than bar 7. After popping it, the stack is empty, so there is no left wall. The note says "bar 3 has no left wall", and the loop stops. Bar 7 is pushed, and the stack is [7].

Look at the totals. Before i = 7, the total was 2. After the basins at 6 and 4 are added, it is 2 + 0 + 3 = 5, as the trace shows.

The final total is 6. The basin at index 2 adds 1 unit, the basin at index 5 adds 1, the basin with bottom 4 adds 3, and the basin at index 9 adds 1. That makes 1 + 1 + 3 + 1 = 6.

Worked example 2: a flat bottom

Flat bottoms are harder, because two bars in a row can have the same height and form one basin with two bottom bars. Take [3, 0, 0, 2, 0, 4]. The water is 10 units, and the program prints 10 for the brute force, the two-pointer scan and the stack.

i  0 h 3 | before []                 | - | total 0
i  1 h 0 | before [0]                | - | total 0
i  2 h 0 | before [0, 1]             | - | total 0
i  3 h 2 | before [0, 1, 2]          | basin at 2 walls 1,3 width 1 level 0 adds 0; basin at 1 walls 0,3 width 2 level 2 adds 4 | total 4
i  4 h 0 | before [0, 3]             | - | total 4
i  5 h 4 | before [0, 3, 4]          | basin at 4 walls 3,5 width 1 level 2 adds 2; basin at 3 walls 0,5 width 4 level 3 adds 4; bar 0 has no left wall | total 10

The two zero bars at indices 1 and 2 are pushed on the list one after the other. When bar 3 of height 2 arrives, it pops bar 2 first. Bar 2's left wall is bar 1, which is also height 0, so the level is 0 and that basin adds nothing. Then bar 1 is popped. Its left wall is bar 0, of height 3, so the basin spans two columns. The level is min(2, 3), which is 2, and the depth is 2 minus 0, which is 2. That adds 2 * 2 = 4.

The stack works here because each popped bar is handled as a slice of the basin. The bottom bars of the same height are popped one at a time, and each slice is counted with its own width and level. Their sum is the right total.

Worked example 3: a deep valley

Take the bars [4, 1, 1, 1, 4]. Three low bars sit between two walls of height 4. Each low bar has water 3 units deep above it, so the water is 3 * 3 = 9 units. The stack method reaches this through several pops. The pops for bars 3 and 2 add 0, because their left walls are no taller than they are. Then bar 1 is popped with the wall at index 0, and the width is 3 and the depth is 3, which adds 9. The program prints valley 9.

A staircase with no valley holds nothing. For [1, 2, 3, 4], every pop finds no left wall, so the total stays 0.

The two-pointer method

The second method reaches the same answer with less memory. Put one index at the left end and one at the right end. Keep two running maximums: the tallest bar seen so far from the left, and the tallest bar seen so far from the right. At each step, look at the two end bars and move the pointer on the lower one.

The lower end is the one to finish. If the left bar is lower than the right bar, the water above the left bar is decided by the tallest bar already seen on the left. The right side has a bar at least as tall as the current right bar, which is taller than the left bar, so the right side cannot be the limit. Add the left running maximum minus the left bar's height. Then move the left pointer one step inward. If the right bar is the lower one, do the same on the right.

def trap_two_pointers(h):
    left, right = 0, len(h) - 1
    left_max = right_max = 0
    water = 0
    while left < right:
        if h[left] < h[right]:
            left_max = max(left_max, h[left])
            water += left_max - h[left]
            left += 1
        else:
            right_max = max(right_max, h[right])
            water += right_max - h[right]
            right -= 1
    return water

The line left, right = 0, len(h) - 1 is a tuple assignment. Python sets both names in one line. The same form is used for the two running maximums: left_max = right_max = 0 sets both to 0.

The order of two lines matters. The running maximum is updated first, and then the water is added. If the bar is a new maximum, the update makes the difference 0, which is correct, because the new bar cannot hold water on its own side. The mistakes section shows what goes wrong when the order is reversed.

The two-pointer method uses O(n) time and O(1) extra space, because it stores only a few numbers. The stack method uses O(n) time and O(n) space, because the stack can hold every bar at once. The three methods gave the same answer on every random test below. The check below runs 3,000 random arrays of up to 14 bars, each with heights from 0 to 7. It uses the three functions from the program above, with the same seed each time you run it, so the output is the same on every run.

import random
 
rng = random.Random(42)
agree = 0
for _ in range(3000):
    n = rng.randint(0, 14)
    arr = [rng.randint(0, 7) for _ in range(n)]
    b = trap_brute_force(arr)
    if trap_with_stack(arr) == b and trap_two_pointers(arr) == b:
        agree += 1
print("random rounds agreeing with brute force:", agree)

random rounds agreeing with brute force: 3000

random.Random(42) makes a random number generator with a fixed seed. randint(0, 14) returns a whole number from 0 to 14, and both ends are included.

Choosing between the two

Use the stack when the task asks for a sequence of basins, such as the position and size of each pool. The stack finds each basin as it closes, so the answer for each basin is available in the order it is found. Use the two-pointer method when only the total is needed and memory matters. Both are standard answers to this problem.

Common mistakes

Using the bar height as the width. The width is the number of bars between the walls, which is the index gap minus 1. Taking the height of a wall gives a number that has nothing to do with the width. For the bars [2, 0, 2], the correct answer is 2. If the line width = i - left - 1 is replaced by width = h[left], the program prints 4 for that input. The correct stack function prints 2.

Forgetting to subtract the bottom's height. The depth of the water is the level minus the height of the bar at the bottom of the basin. Leaving out that subtraction counts the bottom bar itself as water. If the line level = min(height, h[left]) - h[bottom] becomes level = min(height, h[left]), the program prints 12 for the first example instead of 6.

Pushing the current bar at the wrong time, or not at all. The new bar must be pushed after the pops finish, so it can be the left wall for the bars that come later. If stack.append(i) is removed, the stack never has anything to pop, and the water is lost. For the first example, the total is 0.

Updating the running maximum after adding the water. In the two-pointer method, the water for a bar that is a new maximum must be 0. If the water is added before the maximum is updated, the difference is taken against the old maximum. If both branches are reversed, the bars [2, 0, 2] print 0, not 2.

Reading the top of an empty list. After a pop, the list may be empty, which means there is no left wall. In Python, stack[-1] on an empty list does not return None. It raises an error. This is the error from a version of the stack function with the if not stack: break lines removed, run on the first example. The traceback is shortened to the lines that matter:

Traceback (most recent call last):
  File "trap_err_demo.py", line 15, in <module>
    print(trap_no_empty_check([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]))
  File "trap_err_demo.py", line 7, in trap_no_empty_check
    left = stack[-1]
IndexError: list index out of range

The last line gives the error type, IndexError, and the reason, list index out of range. The file name and line numbers depend on where you save the program. The check if not stack: break prevents this.

Confusing the walls with the basin. The water is held between the two walls, so the width covers only the bars between them. Counting the walls as part of the width adds 2 columns that hold no water.

Practice

Exercise 1

Find the total water for [4, 2, 0, 3, 2, 5] by hand, using the brute force idea: for each bar, find the smaller of the tallest bar on its left and the tallest bar on its right, then subtract the bar's own height.

Show the answer

The total is 9. Bar 0 has height 4, and the tallest bar to its right is 5, so the level is 4, and it holds 0. Bar 1 has height 2, the left maximum is 4, and the right maximum is 5, so the level is 4 and it holds 2. Bar 2 has height 0, with the same maxima, so it holds 4. Bar 3 has height 3, and the level is 4, so it holds 1. Bar 4 has height 2, and the level is 4, so it holds 2. Bar 5 has height 5 and holds 0. The sum is 0 + 2 + 4 + 1 + 2 + 0, which is 9. The function trap_brute_force returns 9 for that list.

Exercise 2

Find the total water for [3, 0, 1, 0, 2].

Show the answer

The total is 5. The level above bars 1, 2 and 3 is 2, because the tallest bar on the left is 3 and the tallest bar on the right is 2, and the smaller of the two is 2. Bar 1 holds 2 - 0 = 2. Bar 2 holds 2 - 1 = 1. Bar 3 holds 2 - 0 = 2. The sum is 5. The stack method splits the same water differently. When bar 2 (height 1) arrives, bar 1 is popped with walls 0 and 2, which adds 1. When bar 4 (height 2) arrives, bar 3 is popped with walls 2 and 4, which adds 1, and then bar 2 is popped with walls 0 and 4, which adds 3. The sum is 1 + 1 + 3, which is also 5.

Exercise 3

Explain in plain words why the stack pops only when the new bar is strictly taller than the top bar.

Show the answer

An equal bar does not finish a basin. A basin is finished by a bar that is taller than its bottom, because only then does water sit above the bottom. An equal bar sits at the same height as the bottom, so the water above it would be 0. The condition height > h[stack[-1]] keeps equal bars on the stack, so a flat bottom made of several bars at the same height stays together. When a taller bar arrives, it pops each bar of that group in turn, and each pop is counted with its own width.

Exercise 4

The two-pointer method moves the pointer on the lower end. Explain why the lower end is the one that can be finished.

Show the answer

Suppose the left bar is lower than the right bar. The water above the left bar is the smaller of the tallest bar on its left and the tallest bar on its right, minus the bar's height. The method uses the tallest bar on the left, so it must know that the tallest bar on the right is at least as tall. Here is why. The left pointer moves past a bar only when some bar still to the right of it is taller. So every bar on the left is shorter than some bar on the right, and the tallest bar on the left is shorter than the tallest bar on the right. The left running maximum is therefore the smaller of the two, and the water above the left bar is the running maximum minus the bar's height. The same reasoning works on the right side when the right bar is the lower one.

Exercise 5

Use the first example and stop after the bar at index 7 has been handled. What total has the stack method counted so far, and which bars are on the stack?

Show the answer

The total so far is 5. The trace row for index 7 shows the total after that bar, and the stack after the push is [7], because bars 6, 4 and 3 were all popped. Bar 7 is the tallest bar seen so far, so it is the only bar that is still waiting for a taller bar on its right.

Summary

• Water in a gap rises to the lower of the two walls around it, and its amount is the width times the depth.
• The stack method keeps bars from tallest at the bottom to shortest at the top, and fills one basin each time a taller bar pops its bottom.
• A basin with no left wall holds nothing, so the loop stops there.
• The two-pointer method gives the same total with O(1) extra space, by always finishing the lower side.
• Check the width, the subtraction of the bottom height, the push of the current bar and the order of the running maximum update when a result is wrong.