The Trapped Pool

A hard Python interview practice problem on DataDriven. Write the Python and run it against test cases, with instant feedback.

Domain
Python
Difficulty
hard
Seniority
L5

Problem

A terrain survey hands you an elevation profile as a list of non-negative heights, each bar one unit wide. After a storm, water pools in the dips between taller bars, while anything past the outer edges drains away. Return the total units of water the profile holds once the level settles.

What collects in the valleys after the rain?

Example 1

Input:

heights = [0,1,0,2,1,0,1,3,2,1,2,1]

Output:

6

Example 2

Input:

heights = [4,2,0,3,2,5]

Output:

9

Worked solution and explanation

What this problem really is

Strip the flood imagery and there is one question asked at every index: how high can water stand here? The answer is min(tallest bar to the left, tallest bar to the right) minus this bar's own height, floored at zero. Everyone reaches for two precomputed arrays of running maxima, and that earns a passing grade. What separates seniors is collapsing those arrays into two pointers that store nothing. Miss the insight that the SMALLER of the two running maxima is the binding constraint, and you either overfill the dips or advance the wrong pointer.

Build the reasoning

Step 1: Pin the constraining height at each bar

The water level above bar i is min(max_left[i], max_right[i]) minus height[i], clamped to 0. A bar taller than both walls holds nothing.

Step 2: Replace the arrays with two converging pointers

Track left_max and right_max as pointers move inward. Always step the side whose running max is smaller: that side is the bottleneck, so the water it holds is already known regardless of what the far side does later.

The solution

Two converging pointers tracking running maxima
def trap_rain_water(heights):
    if len(heights) < 3:
        return 0
    left, right = 0, len(heights) - 1
    left_max, right_max = heights[left], heights[right]
    total = 0
    while left < right:
        if left_max <= right_max:
            left += 1
            if heights[left] > left_max:
                left_max = heights[left]
            else:
                total += left_max - heights[left]
        else:
            right -= 1
            if heights[right] > right_max:
                right_max = heights[right]
            else:
                total += right_max - heights[right]
    return total

Common follow-up questions