Calculate how much water a group of walls can hold. Interview Question from @cassidoo's email newsletter.

water_trap.py

"""Calculate how much water is trapped between walls.
                                            |                              |
Basically, wall sections of the pattern |___| are easier to calculate than |___|.  So,
Split the walls up into water-tight sections where the last wall >= the first wall.
If that's not possible, then reverse the whole section so that you have at least one increasing section
and keep working.
   |  |                                                  |  |  |               |
|__|__|__|__|    This gets processed in these chunks: |__|, |__|, |__|, and |__|.  

They hold 1, 2, 1, and 1, respectively, for a total of 5!
"""

import sys


def water_trap(heights):
    """
    Calculate how much water a sequence of walls can hold inside.
    Walls are given as a list of positive integer heights.
    Returns total integer "cells" filled with water
    """
    return sum(score_segment(segment) for segment in parse_segments(heights))


def score_segment(segment):
    """
    Given one single water-tight segment with last wall >= first wall,
    calculate the amount of water inside
    """
    lip_height = segment[0]
    return sum(lip_height - wall_height for wall_height in segment if wall_height < lip_height)


def parse_segments(heights):
    """Returns a number of bins that can be easily calculated for water trapping"""
    walls_remaining = heights[:]
    while len(walls_remaining) > 1:
        next_ind = next_highest_index(walls_remaining)
        if next_ind == 0:
            walls_remaining = list(reversed(walls_remaining))
            continue
        yield walls_remaining[:next_ind+1]
        walls_remaining = walls_remaining[next_ind:]
    raise StopIteration


def next_highest_index(heights):
    """
    Returns the index of the next wall that is higher
    than, or equal to, the current height.
    """
    current = heights[0]

    for ind, height in enumerate(heights[1:], start=1):
        if height >= current:
            return ind

    return 0


# Tests


def test_single_uniform_blank():
    walls = [2, 0, 2]

    # |x|
    # |x|

    assert 2 == water_trap(walls)


def test_fully_submerged_wall():
    walls = [3, 0, 0, 2, 0, 4]

    #      |
    # |xxxx|
    # |xx|x|
    # |xx|x|

    assert 10 == water_trap(walls)


def test_peak_in_middle():
    walls = [2, 0, 0, 4, 0, 0, 2]

    #    |
    #    |
    # |xx|xx|
    # |xx|xx|

    assert 8 == water_trap(walls)


def test_increasing_decreasing_and_equal_segments():
    walls = [1, 0, 4, 0, 1, 0, 4, 0, 1, 0]

    #   |xxx|
    #   |xxx|
    #   |xxx|
    # |x|x|x|x|_

    assert 13 == water_trap(walls)


if __name__ == "__main__":
    walls = [int(height) for height in sys.argv[1:]]
    print(water_trap(walls))