Solution to waterflow problem

gistfile1.py

"""
Single-pass, simple solution to waterflow problem.
"""
q = [2,5,1,3,1,2,1,7,7,6]
def solve(heights):
    water = [0]*len(heights)
    total = 0
    runoff = 0
    for i,height in enumerate(heights):
        if i == 0:
            continue
        above = max(height, heights[i-1] + water[i-1]) - height 
        water[i] = above
        if water[i] == 0:
            runoff = 0
        else:
            runoff += water[i]
    return sum(water)-runoff
print q, solve(q)