DSA Trapping Rain Water

The problem

Trapping Rain Water counts the units of water held between bars after rain. Each bar has a height. Water settles in the dips between taller bars.

A position holds water up to a level set by the walls around it. Water above the two end positions drains away, so the ends themselves hold none.

cases = [
    ([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1], 6),
    ([4, 2, 0, 3, 2, 5], 9),
    ([2, 0, 2], 2),
    ([3, 3, 3], 0),
    ([], 0),
]
for heights, expected in cases:
    print("heights", heights, "-> water", expected)

The third case is the smallest real trap. One dip of depth two sits between two walls of height two. The flat case holds nothing, because no dip exists.

Intuition: each position is capped by its lower wall

A running maximum is the tallest bar seen so far from one side. Each position has a running maximum from the left and one from the right. The smaller of the two caps the water level at that position.

The water at a position is that cap minus the bar height. The formula is simple. The trap is its price when both maxima are rebuilt from scratch for every position.

heights = [3, 1, 2, 0, 2]
scans = 0
total = 0
for i, h in enumerate(heights):
    left_max = 0
    for x in heights[:i + 1]:
        scans += 1
        left_max = max(left_max, x)
    right_max = 0
    for x in heights[i:]:
        scans += 1
        right_max = max(right_max, x)
    total += min(left_max, right_max) - h
print("water:", total)
print("bar visits for maxima:", scans)

Five bars already cost thirty bar visits. Each maximum scan repeats work the neighbouring position just did. The maxima only ever grow as a scan advances, and that fact is wasted here.

Brute force: rebuild both maxima at every position

The brute force writes the intuition as a function. Every position scans to the left end for its left maximum. Then it scans to the right end for its right maximum.

def trap_brute(heights):
    water = 0
    visits = 0
    for i, h in enumerate(heights):
        left_max = 0
        for x in heights[:i + 1]:
            visits += 1
            left_max = max(left_max, x)
        right_max = 0
        for x in heights[i:]:
            visits += 1
            right_max = max(right_max, x)
        water += min(left_max, right_max) - h
    return water, visits

print("main case:", trap_brute([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]))
print("second case:", trap_brute([4, 2, 0, 3, 2, 5]))
print("small case:", trap_brute([2, 0, 2]))

The answers match the cases from the start. The visit counts expose the cost. Twelve bars cost well over a hundred visits, and doubling the bars roughly quadruples that count.

The optimal solution: two pointers carrying their running maxima

Two pointers start at the two ends and carry one running maximum each. The side with the shorter current bar is settled next, because a taller bar on the other side already guarantees its level is capped by its own side.

Settling a position adds its capped level minus its height to the total. Each move settles exactly one position. The pointers meet when every position is settled.

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

cases = [
    ([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1], 6),
    ([4, 2, 0, 3, 2, 5], 9),
    ([2, 0, 2], 2),
    ([3, 3, 3], 0),
    ([], 0),
]
for heights, expected in cases:
    water, moves = trap(heights)
    print("heights", heights, "->", water, "match:", water == expected, "moves:", moves)

All five cases pass. The move count stays below the bar count on every skyline. No table is built, and no position is visited twice.

Walkthrough: [3, 1, 2, 0, 2] step by step

The walk settles one position per line. The state is the two pointers, the two running maxima, and the water total so far.

heights = [3, 1, 2, 0, 2]
left, right = 0, len(heights) - 1
left_max = 0
right_max = 0
water = 0
while left < right:
    if heights[left] < heights[right]:
        left_max = max(left_max, heights[left])
        gained = left_max - heights[left]
        print("settle left index", left, "height", heights[left], "- gained", gained)
        water += gained
        left += 1
    else:
        right_max = max(right_max, heights[right])
        gained = right_max - heights[right]
        print("settle right index", right, "height", heights[right], "- gained", gained)
        water += gained
        right -= 1
print("total water:", water)

The right wall of height 2 settles first and gains nothing. The low bar beside it then gains two units under that cap. The walk settles the middle bar for nothing and the bar of height 1 for one unit, and the total reads three.

Complexity

The brute force scans the whole skyline for each position, so its time is O(n^2) for n bars. The pointer walk settles each position once, so its time is O(n).

The pointer walk stores two maxima, two indexes, and the total, so its extra space is O(1). A precomputed table of left and right maxima would also run in O(n) time, at the price of O(n) space.

def visits_brute(n):
    return sum((i + 1) + (n - i) for i in range(n))

for n in [5, 10, 20]:
    print("bars", n, "- brute visits:", visits_brute(n), " pointer moves:", max(n - 1, 0))

Brute visits grow with the square of the bar count. Pointer moves grow in a straight line. The gap widens with every bar added.

Edge cases and pitfalls

A rising skyline holds no water. Every dip needs a wall on both sides, and a rising run never closes a dip on the right. A falling skyline fails the same way in mirror image.

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

print("empty:", trap([]))
print("one bar:", trap([5]))
print("two bars:", trap([2, 1]))
print("flat:", trap([3, 3, 3]))
print("rising:", trap([1, 2, 3, 4]))
print("falling:", trap([4, 3, 2, 1]))
print("deep single dip:", trap([5, 0, 5]))

The classic pitfall is settling the taller side first. Its cap is not yet known, because a still taller wall may wait beyond the other pointer. The shorter current bar is always the safe side to settle.

Updating the running maximum before adding water matters too. A new tallest bar must add zero for itself. The maximum rises first, and the subtraction then gives zero naturally.