water.py

def create_map(heights):
    max_height = max(heights)
    return '\n'.join(''.join('##' if i >= h else '  ' for i in heights)
                     for h in range(max_height, 0, -1))


def method_a(heights):
    water = 0
    left = 0
    right = len(heights) - 1
    left_max = heights[0]
    right_max = heights[-1]

    #print '[%i<->%i][%i<->%i] -> %i' % (left, right, left_max, right_max, water)
    while left < right:

        if heights[left] > left_max:
            left_max = heights[left]

        if heights[right] > right_max:
            right_max = heights[right]

        if left_max >= right_max:
            water += right_max - heights[right]
            right -= 1
        else:
            water += left_max - heights[left]
            left += 1

        #print '[%i<->%i][%i<->%i] -> %i' % (left, right, left_max, right_max, water)

    return water


def method_b(heights):
    stack = []
    local_max = 0
    water = 0

    for i in heights:
        if i >= local_max:
            water_now = sum([local_max - h for h in stack])
            #print '>>>', water_now, local_max, stack
            local_max = i
            water += water_now
            stack = []
        else:
            stack.append(i)
        #print local_max, stack

    if stack:
        water += method_b(stack[::-1] + [local_max])

    return water


def measure_water(heights, debug=False):
    """
    >>> measure_water([2, 5, 1, 2, 3, 4, 7, 7, 6])
    10
    >>> measure_water([2, 5, 1, 3, 1, 2, 1, 7, 7, 6])
    17
    >>> measure_water([5, 1, 3, 6, 1, 6, 1, 3, 1, 4])
    18
    >>> measure_water([1, 2, 3, 4, 5, 4, 3, 2, 1])
    0
    >>> measure_water([1, 2, 3, 4, 5])
    0
    >>> measure_water([5, 1, 4, 2, 3])
    4
    >>> measure_water([3, 4, 7, 3, 4, 7, 6, 7, 2, 4])
    10
    >>> measure_water([6, 1, 5, 2, 1, 4])
    9
    >>> measure_water([7, 1, 1, 1, 7, 1, 1, 1, 1, 7])
    42
    """

    if debug:
        print create_map(heights)

    return method_a(heights)

if __name__ == '__main__':
    import doctest
    doctest.testmod()

    from time import clock as time

    NUM = 3000

    def test(method):
        assert(method([2, 5, 1, 2, 3, 4, 7, 7, 6]) == 10)
        assert(method([2, 5, 1, 3, 1, 2, 1, 7, 7, 6]) == 17)
        assert(method([5, 1, 3, 6, 1, 6, 1, 3, 1, 4]) == 18)
        assert(method([1, 2, 3, 4, 5, 4, 3, 2, 1]) == 0)
        assert(method([1, 2, 3, 4, 5]) == 0)
        assert(method([5, 1, 4, 2, 3]) == 4)
        assert(method([3, 4, 7, 3, 4, 7, 6, 7, 2, 4]) == 10)
        assert(method([6, 1, 5, 2, 1, 4]) == 9)
        assert(method([7, 1, 1, 1, 7, 1, 1, 1, 1, 7]) == 42)

    count = 10
    num = 50000
    methods = {
        method_a: 0,
        method_b: 0
    }

    for _ in xrange(count):
        for method in methods:
            begin = time()
            for _ in xrange(num):
                test(method)
            elapsed = time() - begin
            methods[method] += elapsed
            print method, elapsed

    print 'Totals:'
    for method, elapsed in methods.iteritems():
        print method, elapsed / count