Algorithm to solve for water accumulation given building heights

Question

I am practicing algorithms, and I have been stuck on this problem for a few days. When I test my solution, I am still incorrect. Here is the problem statement:

The Wall Street in New York is known for its breathtaking skyscrapers. But the raining season is approaching, and the amount of water that will fall over the buildings is going to be huge this year. Since every building is stuck to the buildings to its left and to its right (except for the first and the last one), the water can leak from a building only if the height of the building is higher than the height of the building to its left or to its right (the height of the edges of Wall Street is 0). All the buildings have the width of 1 meter. You are given the heights (in meters) of the buildings on Wall Street from left to right, and your task is to print to the standard output (stdout) the total amount of water (in cubic meters) held over the buildings on Wall Street.

Example input:

heights: [9 8 7 8 9 5 6]

Example output:

5

Explanation: In this example, between the first and the fifth building there are held 4 cubic meters of water (1 over the second, 2 over the third, and 1 over the fourth), and between the fifth and the seventh building there is 1 cubic meter of water held (over the sixth building).

My approach to this problem is to find global maxima, and use differences in these maxima to calculate water accumulation. I account for water I might have missed at the end using the local_water variable. Can anyone help me find the error in my algorithm or code?

NOTE: I am looking for a solution which passes through each element only once

Here is the input where I have an error:

heights: [8,8,4,5]

this input should yield 1, not my answer which is 0.

Here is my code:

def skyscrapers(heights):
    heights.insert(0,0)
    heights.append(0)
    local_max = 0
    global_max = 0
    total_water = 0
    local_water = 0
    end_water = []
        # end_water records water heights to be used for finding 
                # water between the final global maximum and 
                # subsequent local maximums. These potential values are
                # stored in local_water.
    for i in range(1, len(heights)-1):
        end_water.append(heights[i]) 

        # check for local max
        if heights[i-1] < heights[i] and heights[i] > heights[i+1]:

            # record potential collected water for after final
            # gloabl max
            for s in end_water:
                local_water += (heights[i] - s) * (heights[i] - s > 0)
            local_max=i

        # new global max
        if heights[local_max] > heights[global_max]:
            for s in heights[global_max:local_max]:
                if heights[global_max] - s > 0:
                    total_water += heights[global_max] - s
            global_max = local_max
            local_water = 0
            end_water = []

    total_water += local_water

    print total_water

Answer 1

height
   _       _
9 | |_   _| |      _ _
8 |   |_|   |     |   |
7 |         |  _  |   |
6 |         |_| | |   |  _
5 |             | |   |_| |
4 |             | |       |  _       _ 
3 |             | |       | | |  _  | |
2 |             | |       | | |_| |_| |
1 |0 1 2 3 4 5 6| |0 1 2 3| |0 1 2 3 4| pos

Here's a single-pass (!) (O(n)-time) O(n)-space algorithm based on the stack-based solution for Maximize the rectangular area under Histogram problem:

from collections import namedtuple

Wall = namedtuple('Wall', 'pos height')

def max_water_heldover(heights):
    """Find the maximum amount of water held over skyscrapers with *heights*."""
    stack = []
    water_held = 0 # the total amount of water held over skyscrapers
    for pos, height in enumerate(heights):
        while True:
            if not stack or height < stack[-1].height: # downhill
                stack.append(Wall(pos + 1, height)) # push possible left well wall
                break
            else: # height >= stack[-1].height -- found the right well wall/bottom
                bottom = stack.pop().height
                if stack: # if there is the left well wall
                    well_height = min(stack[-1].height, height) - bottom
                    if well_height:
                        water_held += (pos - stack[-1].pos) * well_height
    return water_held

Example:

>>> max_water_heldover([9, 8, 7, 8, 9, 5, 6])
5
>>> max_water_heldover([8, 8, 4, 5])
1
>>> max_water_heldover([3, 1, 2, 1, 3])
5

I've tested it against a brute-force O(n*m) algorithm:

from itertools import product

def test(max_buildings=6, max_floors=6):
    for nbuildings, nfloors in product(range(max_buildings), range(max_floors)):
        for heights in product(range(nfloors), repeat=nbuildings):
            assert max_water_heldover(heights) == max_water_heldover_bruteforce(heights), heights

where max_water_heldover_bruteforce() is:

import sys
from colorama import Back, Fore, Style, init # $ pip install colorama
init(strip=not sys.stderr.isatty())

def max_water_heldover_bruteforce(heights):
    if not heights: return 0
    BUILDING, AIR, WATER = ['B', ' ',
            Back.BLUE + Fore.WHITE + 'W' + Fore.RESET + Back.RESET + Style.RESET_ALL]
    matrix = [[WATER] * len(heights) for _ in range(max(heights))]
    for current_floor, skyscrapers in enumerate(matrix, start=1):
        outside = True
        for building_number, building_height in enumerate(heights):
            if current_floor <= building_height:
                outside = False
                skyscrapers[building_number] = BUILDING
            elif outside:
                skyscrapers[building_number] = AIR
        for i, building_height in enumerate(reversed(heights), 1):
            if current_floor > building_height:
                skyscrapers[-i] = AIR
            else:
                break
    if '--draw-skyscrapers' in sys.argv:
        print('\n'.join(map(''.join, matrix[::-1])), file=sys.stderr)
        print('-'*60, file=sys.stderr)
    return sum(1 for row in matrix for cell in row if cell == WATER)

The results are the same.