python - 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

Here's a one-pass solution that improves on liuzhidong's and J.S. Sebastian's solutions by using only O(1) space:

def fillcount(elevations):
    start = 0
    end = len(elevations) - 1
    count = 0
    leftmax = 0
    rightmax = 0

    while start < end:
        while start < end and elevations[start] <= elevations[start + 1]:
            start += 1
        else:
            leftmax = elevations[start]

        while end > start and elevations[end] <= elevations[end - 1]:
            end -= 1
        else:
            rightmax = elevations[end]

        if leftmax < rightmax:
            start += 1
            while start < end and elevations[start] <= leftmax:
                count += leftmax - elevations[start]
                start += 1
        else:
            end -= 1
            while end > start and elevations[end] <= rightmax:
                count += rightmax - elevations[end]
                end -= 1

    return count

I tested it against this simpler two-pass solution:

def fillcount_twopass(elevations):
    global_max = max(range(len(elevations)), key=lambda x: elevations[x])
    count = 0
    local_max = 0

    for i in xrange(0, global_max):
        if elevations[i] > local_max:
            local_max = elevations[i]
        else:
            count += local_max - elevations[i]

    local_max = 0
    for i in xrange(len(elevations) - 1, global_max, -1):
        if elevations[i] > local_max:
            local_max = elevations[i]
        else:
            count += local_max - elevations[i]

    return count

The two-pass solution is based on the following logic:

1. Find the maximum peak for the whole graph -- the global maximum.
2. Scan from the left side towards the global maximum peak. Keep track of the left maximum. Every cell that is a) below or at the same level as the left maximum, b) to the right of the left maximum, and c) to the left of the global maximum will hold water. When the left maximum increases, it has no effect on the earlier columns, but the later columns now hold water at this new maximum level.
3. Do the same from the right, in reverse.

This is similar to what Rémi suggested, but uses the global maximum to provide an anchor, which simplifies things.

The one-pass solution is partially based on ideas from Mark Tolonen. It improves on the above by observing that we can do both the left and right pass simultaneously, without knowing the global maximum. That's because the current maximum on either side is either greater than, lower than, or equal to the maximum on the other side. The lower maximum will always be lower than the global maximum, even if we don't yet know what the global maximum is, so we can proceed from there to the next local maximum on that side. The algorithm in full detail:

1. Start with pointers at the start and the end of the list, and initialize left_max, right_max, and count to 0.
2. Scan right, incrementing start until you reach a left maximum. Then scan left, decrementing end until you reach a right maximum.
3. From the lower maximum, continue scanning until you reach a column greater than the local maximum, counting fillable cells along the way and adding them to count.
4. Repeat steps 2 and 3, ending when start and end coincide.

Note that for our purposes, a local maximum is simply any point that is preceded by an ascent (and possibly a plateau), and followed by a descent. Local maxima below the highest local maximum encountered so far are only encountered in step 3, where they have no effect.

This last solution can process ten million datapoints in 3 seconds:

>>> rands = [random.randrange(0, 1000000) for i in xrange(10000000)]
>>> %timeit fillcount(rands)
1 loops, best of 3: 3.3 s per loop