Trapping rain water with a stack

Why this matters

Trapping rain water is one of the best-known stack problems in interviews. The question is simple to picture. Rain falls on a row of blocks. Some water sits in the dips between taller blocks and stays there. How much water is trapped? The problem is useful because two very different methods solve it, and a stack method teaches you how to think about nested "walls" that hide smaller "basins".

This lesson builds the stack method from a picture, one water level at a time. It then shows the two-pointer method, which solves the same problem without a stack, and a column-by-column loop that is easy to check. You will run all three, compare their answers, and see the mistakes that crash or quietly give wrong numbers.

The picture

Each block has a height. Use a list of whole numbers, where each number is the height of one block. Each block is one unit wide. Here is the list 0 1 0 2 1 0 1 3 2 1 2 1, drawn with # for a block and ~ for trapped water:

              #
      # ~ ~ ~ # # ~ #
  # ~ # # ~ # # # # # #
0 1 0 2 1 0 1 3 2 1 2 1

The top row shows the tallest block, which has height 3. Each row below is one level higher than the row above it. The water in the dips is shown as ~. Count the ~ marks and you get 6 units of water. A unit is one block-wide, one block-high cell of water.

Notice what holds the water. A dip holds water only if there is a wall on its left and a wall on its right, and both walls are taller than the dip. The water sits up to the level of the shorter wall. Water in the dip at column 5, which has height 0, rises to level 2. It is held on the left by the block at column 3 (height 2) and on the right by the block at column 7 (height 3). The shorter of those walls is 2, so the water stops at 2. Water spills over the lower wall, so it never gets higher than that.

A block at the very end, with no wall on one side, holds no water. Water runs off the open side. The list 5 4 3 2 1 has no water at all, because nothing stops the water from running off to the right.

The water over one column

Start with a simple rule for one column. A column is the position of one block. The water above that column is the level of the shorter of two walls, one on the left and one on the right, minus the column's own height. The walls are the tallest blocks to the left and to the right of the column, counting the column itself.

Here is a worked example for column 5 in the list 0 1 0 2 1 0 1 3 2 1 2 1. The column height is 0. The tallest block to the left of it, including column 5, is 2 (column 3). The tallest block to the right, including column 5, is 3 (column 7). The shorter wall is 2. The water is 2 minus 0, which is 2 units. That is two of the six units in the picture.

Doing this for each column is correct, and it is the easiest version to check. The code is short:

int trap_columns(const std::vector<int>& h) {
    int water = 0;
    for (std::size_t i = 0; i < h.size(); ++i) {
        int left = 0;
        int right = 0;
        for (std::size_t j = 0; j <= i; ++j) left = std::max(left, h[j]);
        for (std::size_t j = i; j < h.size(); ++j) right = std::max(right, h[j]);
        water += std::min(left, right) - h[i];
    }
    return water;
}

The two inner loops scan the whole list for every column. Each column costs about n steps, and there are n columns, so the work is O(n squared). That is fine for a short list and too slow for a list with a hundred thousand blocks. The stack method below does the same job in one pass.

The stack idea: fill one layer at a time

Instead of working column by column, the stack method fills the water in horizontal layers. Each layer is a slab of water between two walls and one level. When a bar is taller than the bar on top of the stack, the bar on top is sitting in a basin. The walls of that basin are the bar below it in the stack (on the left) and the new bar (on the right). The water in the basin rises to the level of the shorter wall, and the slab above the basin's bottom bar is filled.

The stack holds the positions of bars that might still hold water on their right. It is kept in decreasing order from bottom to top. The rule is this. For each new bar at position i:

1. While the stack is not empty and the new bar is taller than the bar on top, the top bar is at the bottom of a basin. Pop it. Call it bottom.
2. If the stack is now empty, there is no wall on the left. Stop, because no water can be held here.
3. Otherwise the new top of the stack is the left wall, at position left. The basin is width columns wide, where width = i - left - 1. Its water level is min(h[left], h[i]), the shorter of the two walls. The slab adds width times (level minus h[bottom]) units of water.
4. Push i.

The idea that makes this correct is in step 3. The popped bar bottom is the tallest bar between the two walls. Every bar in the gap is at or below it. So the slab of water between h[bottom] and the level sits above every bar in the gap. Its width is the full gap, because all the bars under it are below its floor.

The stack method on one example

Take the list 4 2 0 3 2 5. It is a good first example, because one bar traps water at two different levels. Here is the picture, with the same symbols as before:

          #
# ~ ~ ~ ~ #
# ~ ~ # ~ #
# # ~ # # #
# # ~ # # #
4 2 0 3 2 5

The water is 9 units. Here is the trace, one row per new bar. The "walls" column shows the stack after the step, with the bar heights in brackets.

Check the last row against the picture. Step 5 adds two slabs. Popping bar 4 adds a slab one column wide at level 3, which fills column 4. Popping bar 3 adds a slab four columns wide (positions 1 to 4) at level 4, which gives one unit to each of those columns. Step 3 had already added two units to column 2 (levels 1 and 2) and one unit each to columns 1 and 2 at level 3. Add up each column: column 1 holds 2, column 2 holds 4, column 3 holds 1 and column 4 holds 2. The total is 9.

The code for the method is below. It is the same rule as the trace, written in C++.

int trap_stack(const std::vector<int>& h) {
    int water = 0;
    std::stack<std::size_t> walls;  // bars that may still hold water on their right
    for (std::size_t i = 0; i < h.size(); ++i) {
        while (!walls.empty() && h[i] > h[walls.top()]) {
            std::size_t bottom = walls.top();
            walls.pop();
            if (walls.empty()) break;  // no wall on the left, so no water here
            std::size_t left = walls.top();
            std::size_t width = i - left - 1;
            int level = std::min(h[left], h[i]);
            water += static_cast<int>(width) * (level - h[bottom]);
        }
        walls.push(i);
    }
    return water;
}

Two lines need attention. The if (walls.empty()) break; line handles the case where the popped bar is the first bar in the list, with no wall on the left. The width line counts only the columns between the two walls, so it uses i - left - 1. The static_cast<int>(width) turns the unsigned width into an int before the multiplication, so the sum stays signed.

The two-pointer method

There is a second method that uses no stack at all. It walks two pointers inward, one from each end of the list. It keeps the tallest wall seen so far from each side. The rule is simple: move the pointer on the side whose wall is shorter. The water on that side is set by its own tallest wall, because the other side is taller and cannot be the limit there.

int trap_two_pointers(const std::vector<int>& h) {
    if (h.empty()) return 0;
    std::size_t lo = 0;
    std::size_t hi = h.size() - 1;
    int left_max = 0;
    int right_max = 0;
    int water = 0;
    while (lo < hi) {
        if (h[lo] < h[hi]) {
            left_max = std::max(left_max, h[lo]);
            water += left_max - h[lo];
            ++lo;
        } else {
            right_max = std::max(right_max, h[hi]);
            water += right_max - h[hi];
            --hi;
        }
    }
    return water;
}

Trace it on 4 2 0 3 2 5. The pointers start at lo = 0 (height 4) and hi = 5 (height 5).

At step 2, the left side moved, and the left_max was already 4 from step 1. The water at position 1 (height 2) is 4 minus 2, which is 2. That matches the 2 units the stack method put in column 1, which came from its level 3 and level 4 slabs. The two methods reach the same total by different routes.

Both methods run in one pass and use little extra memory. The two-pointer version uses two integers. The stack version uses a stack that can hold up to n positions. The two-pointer version is shorter, but it is harder to explain why the choice of side is safe. The stack version is easier to follow as a story of water levels, which is why this lesson teaches it first.

A harder example: walls with a wide basin

The first example had one basin with two levels. A harder example has a wide basin. Take 3 0 0 0 3. Three dips of height 0 sit between two walls of height 3. The water rises to level 3 across all three dips, so the answer is 9.

Here is what the stack does. Bar 0 (height 3) is pushed. Bars 1, 2 and 3 (all height 0) are pushed, because none of them is taller than the bar on top. Bar 4 (height 3) arrives. It is taller than bar 3 (height 0), so bar 3 is popped, with left wall bar 2 at height 0. Its width is 4 minus 2 minus 1, which is 1. Its level is min(0, 3), which is 0, so the slab adds 0 units. Then bar 2 is popped the same way, adding 0 units. Then bar 1 is popped with left wall bar 0 (height 3). Its width is 4 minus 0 minus 1, which is 3. Its level is min(3, 3), which is 3. The slab adds 3 times (3 minus 0), which is 9.

The last pop is where the work happens. A stack trace that stops at the first pop would miss the water. Read the last pop as a single wide slab: three columns at level 3. This is the case that the column-by-column loop handles with three separate calculations, and the stack handles with one.

Run the three functions on this and on the other lists in the test program:

A: stack=6 two_pointers=6 columns=6
B: stack=9 two_pointers=9 columns=9
C: stack=2 two_pointers=2 columns=2
D: stack=0 two_pointers=0 columns=0
E: stack=0 two_pointers=0 columns=0
random mismatches: 0

Here A is 0 1 0 2 1 0 1 3 2 1 2 1, B is 4 2 0 3 2 5, C is 2 0 2, D is 5 4 3 2 1, and E is an empty list. The random test built 3000 lists with 0 to 25 blocks, each block of height 0 to 6, and the three methods agreed on every list. A small random test like this one is the best check for a method you have just written.

The stack version gives the water for 3 0 0 0 3 as 9, and the two-pointer version gives the same, so the wide basin is handled by both. The first bar that could go wrong is the one that decides the width, so it is worth tracing one wide case by hand before trusting the code.

Common mistakes

Using the height of the new bar instead of the shorter wall. A common mistake is to write int level = h[i];, which uses the right wall alone. For 4 2 0 3 2 5 that gives a total of 16 instead of 9. The water must stop at the shorter wall, because the taller wall does not hold it in. Real output from the test program with that one line changed:

correct=9 level=h[i]:16 width=i-left:14

Forgetting to subtract one from the width. The gap between two walls has i - left - 1 columns. Without the - 1, the code counts the left wall as part of the gap. For the same list the total becomes 14, which is also wrong. The same output line shows it, in the width=i-left column.

Calling top() or pop() on an empty stack. The test if (walls.empty()) break; is needed. Without it, the code reads walls.top() after the stack is empty. Take the list 1 2. Bar 1 (height 2) pops bar 0, the stack is empty, and the next line tries to read the left wall. With the standard library's checks turned on (-D_GLIBCXX_ASSERTIONS), the program stops with a message like this one and exit code 134. The path and line number depend on the compiler version:

/usr/include/c++/15/bits/stl_stack.h:260: std::stack<_Tp, _Sequence>::reference std::stack<_Tp, _Sequence>::top() [with _Tp = long unsigned int; _Sequence = std::deque<long unsigned int, std::allocator<long unsigned int> >; reference = long unsigned int&]: Assertion '!this->empty()' failed.
exit=134

Without the checks, the same code reads memory it does not own. It may crash, or it may give a wrong total. The fix is the empty() check before the read.

Changing the pop test to >=. Changing > to >= in the stack method still gives the right totals. An equal bar pops with a slab of height 0, so nothing is added. The test program checked this on 3000 random lists with no mismatch. The walkthrough in this lesson uses >, so keep that version when you explain the code.

Moving the taller pointer in the two-pointer method. The two-pointer method works only because it moves the side with the shorter wall. If you move the taller side, the water on the shorter side is counted with the wrong maximum. Changing the test to h[lo] > h[hi] gives 7 on the list 4 2 0 3 2 5, where the correct total is 9. The same change was wrong on most of 3000 random lists. Keep the test h[lo] < h[hi] exactly as written.

Practice

Exercise 1. Count the water over the list 1 0 2 0 1 using the column-by-column method. Give the water in each column and the total.

Show the answer

Column 0 (height 1) holds no water, because it is at the left edge. Column 1 (height 0) has walls 1 on the left and 2 on the right, so the level is 1. Water is 1 minus 0, which is 1. Column 2 (height 2) is a wall, so it holds 0. Column 3 (height 0) has walls 2 on the left and 1 on the right, so the level is 1. Water is 1. Column 4 holds 0. The total is 2. The test program prints 2 for this list with all three methods.

Exercise 2. Trace trap_stack on 2 0 2. Write the stack after each bar and the total after each step.

Show the answer

Bar 0 (height 2) is pushed. Bar 1 (height 0) is not taller than the top, so it is pushed. Bar 2 (height 2) is taller than the top (height 0). Pop bar 1. The stack still has bar 0, so the left wall is 0. Width is 2 minus 0 minus 1, which is 1. Level is min(2, 2), which is 2. The slab adds 1 times (2 minus 0), which is 2. The total is 2. Bar 2 (height 2) is not taller than bar 0 (height 2), so the loop stops. Bar 2 is pushed. The final stack is bars 0 and 2, and the total is 2.

Exercise 3. Count how many columns hold water in the list 4 2 0 3 2 5.

Show the answer

Use the column method. Column 0 holds 0 (it is the left edge). Column 1 (height 2) has a level of min(4, 5) = 4, so it holds 2. Column 2 (height 0) holds 4. Column 3 (height 3) holds min(4, 5) minus 3, which is 1. Column 4 (height 2) holds min(4, 5) minus 2, which is 2. Column 5 (the right edge) holds 0. Four columns hold water: columns 1, 2, 3 and 4. The total is 2 + 4 + 1 + 2, which is 9.

Exercise 4. Explain why the two-pointer method can move the lower side and never needs to store the water of past columns.

Show the answer

The lower side has a wall that is shorter than the other side's largest wall. So the water above the current position on the lower side is set by the lower side's own tallest wall, which the pointer has already passed. The other side is at least as tall, so it cannot be the limit. The water for that column can be added the moment the pointer reaches it, so nothing needs to be stored for later.

Exercise 5. Rewrite trap_stack so it uses a std::vector<std::size_t> as the stack instead of std::stack.

Show the answer

Replace std::stack<std::size_t> walls; with std::vector<std::size_t> walls;. Replace walls.top() with walls.back(), walls.pop() with walls.pop_back(), and walls.push(i) with walls.push_back(i). Keep the empty() checks as they are. The std::vector version gives the same totals on every list. Use std::stack when you want the code to say "stack" and to stop anyone from reading below the top. Use std::vector when you also need to look at other entries.

Summary

• Trapped water is held between two walls, and the water level is the shorter of the two.
• The stack method fills the water in horizontal layers, one popped basin at a time, and its width is the gap between two walls.
• The column-by-column method is slow but easy to check, so use it to test the faster methods on random lists.
• The two-pointer method reaches the same totals in one pass with two integers, by always moving the shorter side.
• The common bugs are using the wrong wall as the level, counting the left wall in the width, and reading the stack when it is empty.