How much water is collected between towers with different heights

water_tower.py

import numpy as np
from itertools import izip

def _fold(func, iterable, iterize):
    _iter = iterize(iterable)
    res = _iter.next()
    yield res
    for i in _iter:
        res = func(res, i)
        yield res

def foldl(func, iterable):
    for i in _fold(func, iterable, iter):
        yield i

def foldr(func, iterable):
    for i in _fold(func, iterable, reversed):
        yield i

def water_capacity(a):
    # highest towers to the left and to the right of each individual tower
    max_lefts = foldl(max, a)
    max_rights = reversed(list(foldr(max, a)))
    
    # the height of water trapped in each tower
    heights = ( min(x,y) for x, y in izip(max_lefts, max_rights) )
    
    # the amount of water is just the height minus the tower height
    return sum( x-y if x > y else 0 for x, y in izip(heights, a))

if __name__ == '__main__':
    a = np.array([1, 5, 2, 3, 5, 2, 3])
    print water_capacity(a)