water.py

def water(bars):
    # get the volume of water between two peaks.
    def volume_between_peaks(i, j):
        # first make all the bars between the two peaks lower than the lower peak.
        if bars[i] > bars[j]:
            while bars[i] > bars[j]:
                i += 1
            i -= 1
        else:
            while bars[i] <= bars[j]:
                j -= 1
            j += 1

        # then calculate the volume
        square = (j- i - 1) * min(bars[i], bars[j])
        return square - sum(bars[i+1:j])

    # calculate total volume.
    def volume(peaks):
        if len(peaks) < 2: return 0
        sorted_peaks = sorted(peaks, key=lambda peak: -bars[peak])
        p1, p2 = sorted([sorted_peaks[0], sorted_peaks[1]])
        return sum((volume_between_peaks(p1, p2),
                    volume([p for p in peaks if p <= p1]),
                    volume([p for p in peaks if p >= p2])))

    # get all peak bars
    peaks = []
    _bars = [0] + bars + [0]
    for i in range(1, len(bars)+1):
        if _bars[i-1] < _bars[i] >= _bars[i+1]:
            peaks.append(i-1)
    return volume(peaks)


if __name__ == "__main__":
    testcases = [
        ([2,0,1], 1),
        ([1,0,1], 1),
        ([5,0,5], 5),
        ([0,1,0,1,0], 1),
        ([1,0,1,0], 1),
        ([1,0,1,2,0,2], 3),
        ([2,5,1,2,3,4,7,7,6], 10),
        ([5,1,0,1],1),                 # thanks https://news.ycombinator.com/item?id=6640085
        ([2,5,1,2,3,4,7,7,6,3,5], 12), # thanks https://news.ycombinator.com/item?id=6640105
    ]

    for case in testcases:
      w = water(case[0])
      if w == case[1]:
         print "TRUE: %s holds %s" % (case[0], w)
      else:
         print "MISMATCH: %s holds %s (got %s)" % (case[0], case[1], w)