task.py

import unittest


class MaybeCls:
    def __repr__(self):
        return 'Maybe'
Maybe = MaybeCls()


def pour(wall):

    print wall
    
    def is_hollow(i):
        # Edge of the wall.
        if i == 0 or i == len(wall) - 1:
            return False
        # Isn't a hollow.
        elif wall[i - 1] < wall[i] and wall[i + 1] < wall[i]:
            return False
        # It is.
        elif wall[i - 1] > wall[i] and wall[i + 1] > wall[i]:
            return True
        # Maybe at the bottom of a bigger hollow.
        elif wall[i - 1] == wall[i] or wall[i + 1] == wall[i]:
            return Maybe
        else:
            # Ladder near a hollow (e.g. 2 in [0123])
            return False

    def found_bigger_hollow(start_idx):
        left_boundary = right_boundary = None
        # Search left.
        for i in range(start_idx - 1, -1, -1):
            if wall[i] > wall[start_idx]:
                left_boundary = i
                break
        # Search right.
        for i in range(start_idx + 1, len(wall)):
            if wall[i] > wall[start_idx]:
                right_boundary = i
                break
        if left_boundary is not None and right_boundary is not None:
            return left_boundary, right_boundary
        return ()

    def pour_bigger_hollow(left_boundary, right_boundary):
        print 'pour_bigger_hollow, [%d, %d]' % (left_boundary, right_boundary)
        poured = 0
        height = min(wall[left_boundary], wall[right_boundary])
        for i in range(left_boundary + 1, right_boundary):
            poured += height - wall[i]
            wall[i] = height
        return poured

    for i, level in enumerate(wall):
        hollow = is_hollow(i)
        print i, hollow
        if not hollow:
            continue
        elif hollow is True:
            # Pour the hollow.
            wall[i] += 1
            # Run itself on modified wall.
            new_volume = pour(wall)
            # Return new volume + one poured unit.
            return new_volume + 1
        elif hollow is Maybe:
            bigger_hollow = found_bigger_hollow(i)
            if not bigger_hollow:
                continue
            poured = pour_bigger_hollow(*bigger_hollow)
            new_volume = pour(wall)
            return new_volume + poured
        else:
            raise AssertionError('Unreachable')

    # Couldn't found anything to pour.
    return 0


def found_puddle_volume(wall):
    return pour(wall)


class Tests(unittest.TestCase):

    def test_ac_1(self):
        wall = [2, 5, 1, 2, 3, 4, 7, 7, 6]
        res = found_puddle_volume(wall)
        self.assertEqual(res, 10)

    def test_ac_2(self):
        wall = [2, 5, 1, 3, 1, 2, 1, 7, 7, 6]
        res = found_puddle_volume(wall)
        self.assertEqual(res, 17)

    def test_0(self):
        wall = [2, 1, 2]
        res = found_puddle_volume(wall)
        self.assertEqual(res, 1)

    def test_1(self):
        wall = [2, 1, 1, 2]
        res = found_puddle_volume(wall)
        self.assertEqual(res, 2)

    def test_2(self):
        wall = [3, 1, 2, 3]
        res = found_puddle_volume(wall)
        self.assertEqual(res, 3)

    def test_3(self):
        wall = [3, 1, 2, 3, 1, 1, 8]
        res = found_puddle_volume(wall)
        self.assertEqual(res, 7)

    def test_4(self):
        wall = [2, 1, 4, 3, 6, 3, 4]
        res = found_puddle_volume(wall)
        self.assertEqual(res, 3)


if __name__ == '__main__':
    unittest.main()