The Water Collector

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
Asked in
Sigmoid

Problem

A survey gives you the heights of a row of vertical walls. Any two walls, together with the ground between them, form a container whose capacity is the height of the shorter wall times the distance between the two. Return the greatest capacity any pair of walls can hold.

Two walls, one sky, and a very important question.

Example

Input:

heights = [1,8,6,2,5,4,8,3,7]

Output:

49

Worked solution and explanation

The question hiding under the water

Beneath the wall-and-water costume this is a width-versus-height tradeoff, and the whole problem collapses to one question: when you shrink the width by one step, which wall do you give up? Brute force compares all n^2 pairs and everyone writes it in thirty seconds. On a million walls it never returns. The move that separates candidates is starting at the widest container and always retreating the SHORTER wall. Retreat the taller one instead and you throw away width for a height you were never allowed to reach, and you can silently miss the true maximum.

Building the pass

Step 1: Anchor a pointer at each end

Left at index 0, right at the last index. This is the maximum possible width, so it is the right place to start before width can only shrink.

Step 2: Score the current container

Capacity is the shorter wall's height times the gap between the pointers. Track the best seen so far in a single running variable.

Step 3: Retreat the shorter wall

The shorter wall is the bottleneck, and it can never do better than its current width, so retire it. The taller wall stays: a future partner might finally exceed the current ceiling.

The solution

Two converging pointers, always retire the shorter wall
def max_water(heights):
    left = 0
    right = len(heights) - 1
    best = 0
    while left < right:
        height = min(heights[left], heights[right])
        best = max(best, height * (right - left))
        if heights[left] <= heights[right]:
            left += 1
        else:
            right -= 1
    return best

Common follow-up questions