Water Tanks Problem

watertanks.py

# gets list of key walls from a sorted list of walls
def _keywalls(walls):
    prevheight = pos = highest = 0
    keywalls = []
    for pos, height in walls:
        if height > highest: # raises globally
            highest = height
        elif height >= prevheight: # raises locally
            while keywalls[-1][1] < height:
                keywalls.pop() # discard recent walls until we see with one higher than height (always occurs)
        prevheight = height
        keywalls.append((pos, height))
    return keywalls

# gets total volume given list of key walls
def _watervolume(keywalls):
    prevpos = prevheight = vol = 0
    for pos, height in keywalls:
        level = prevheight if height > prevheight else height
        vol += level * (pos - prevpos)
        prevpos, prevheight = pos, height
    return vol

def watertanks(heights):
    return _watervolume(_keywalls([walls for walls in enumerate(heights)]))

assert(watertanks([4, 2, 1, 3]) == 9)
assert(watertanks([1, 2, 3]) == 3)
assert(watertanks([3, 2, 1]) == 3)
assert(watertanks([1, 2, 3, 2, 1]) == 6)
assert(watertanks([3, 2, 1, 2, 3]) == 12)
assert(watertanks([1, 2, 3, 2, 1, 2, 3, 2, 1]) == 18)
assert(watertanks([3, 2, 1, 2, 3, 2, 1, 2, 3]) == 24)