test.py

import unittest

from water import water, local_min, solve


class LocalMinTest(unittest.TestCase):
    def test_empty(self):
        self.assertEqual(local_min([]), [])

    def test_single(self):
        self.assertEqual(local_min([0]), [])
        self.assertEqual(local_min([1]), [])

    def test_double(self):
        self.assertEqual(local_min([0, 0]), [])
        self.assertEqual(local_min([0, 1]), [])
        self.assertEqual(local_min([1, 0]), [])
        self.assertEqual(local_min([1, 1]), [])

    def test_three(self):
        self.assertEqual(local_min([0, 0, 0]), [])
        self.assertEqual(local_min([0, 0, 1]), [])
        self.assertEqual(local_min([0, 1, 0]), [])
        self.assertEqual(local_min([0, 1, 1]), [])
        self.assertEqual(local_min([1, 0, 0]), [])
        self.assertEqual(local_min([1, 0, 1]), [1])
        self.assertEqual(local_min([1, 1, 0]), [])
        self.assertEqual(local_min([1, 1, 1]), [])

    def test_wide(self):
        self.assertEqual(local_min([1, 0, 0, 1]), [1, 2])

    def test_multiple(self):
        self.assertEqual(local_min([1, 0, 1, 0, 1]), [1, 3])

    def test_right_slope(self):
        self.assertEqual(local_min([2, 1, 0, 1]), [2])


class WaterTest(unittest.TestCase):

    def test_empty(self):
        self.assertEqual(water([]), ([], 0))

    def test_one(self):
        self.assertEqual(water([1]), ([1], 0))

    def test_bucket(self):
        self.assertEqual(water([1, 0, 1]), ([1, 1, 1], 1))

    def test_flat(self):
        self.assertEqual(water([1, 1, 1]), ([1, 1, 1], 0))
        self.assertEqual(water([0, 0, 0]), ([0, 0, 0], 0))

    def test_no_edge(self):
        self.assertEqual(water([0, 0, 1]), ([0, 0, 1], 0))
        self.assertEqual(water([1, 0, 0]), ([1, 0, 0], 0))

    def test_wide(self):
        self.assertEqual(water([1, 0, 0, 1]), ([1, 1, 1, 1], 2))

    def test_multiple(self):
        self.assertEqual(water([1, 0, 1, 0, 1]), ([1, 1, 1, 1, 1], 2))

    def test_right_slope(self):
        self.assertEqual(water([2, 1, 0, 1]), ([2, 1, 1, 1], 1))


class MainTest(unittest.TestCase):
    def test_empty(self):
        self.assertEqual(solve([]), 0)

    def test_one(self):
        self.assertEqual(solve([0]), 0)

    def test_two(self):
        self.assertEqual(solve([0, 0]), 0)

    def test_bowl(self):
        self.assertEqual(solve([1, 0, 1]), 1)

    def test_deep_bowl(self):
        self.assertEqual(solve([2, 0, 2]), 2)

    def test_real(self):
        self.assertEqual(solve([2, 5, 1, 2, 3, 4, 7, 7, 6]), 10)

    def test_right_slope(self):
        self.assertEqual(solve([2, 1, 0, 1]), 1)

water.py

def _edge(arg, start, func):
    return next(x for x in range(start+1, len(arg)) if func(arg[x-1], arg[x]))


def local_min(arg):
    result = []
    position = 0
    try:
        while True:
            position = _edge(arg, position, lambda a, b: a > b)
            i = _edge(arg, position, lambda a, b: a != b)
            if arg[i-1] < arg[i]:
                result.extend(range(position, i))
            else:
                i -= 1
            position = i
    except StopIteration:
        pass
    return result


def water(arg):
    l_m = local_min(arg)
    for i in l_m:
        arg[i] += 1
    return arg, len(l_m)


def solve(arg):
    result = 0
    while True:
        w = water(arg)[0]
        if w == 0:
            break
        result += w
    return result