triangles.py

import re
import sys


_RE_NUM = re.compile(r'\d+')


def get_triangle_path(stream, key=lambda x: x):
    col = 0 #the column number in which the last step of the path was found

    for line in stream:
        opts = map(int, _RE_NUM.findall(line)[col:col+2])

        if len(opts) == 2:
            a, b = opts
            if key(a) >= key(b):
                yield col, a
            else:
                col += 1
                yield col, b

        elif len(opts) == 1:
            yield col, opts[0]


def get_triangle_values(stream, key=lambda x: x):
    return (x[1] for x in get_triangle_path(stream, key))



### TESTS ###

import unittest
from StringIO import StringIO

class TestTrianglePath(unittest.TestCase):
    TR_SMALL = """
            5
           9 6
          4 6 8
         0 7 1 5
    """

    TR_WS = """
        5
               9         	6
      4 6    8

            0  7           1  5
    """

    TR_DUPES = """
            4
           7 6
          3 3 3
         0 1 2 5
    """

    def test_path_small_max(self):
        path = tuple(get_triangle_path(StringIO(self.TR_SMALL)))
        assert path == ((0, 5), (0, 9), (1, 6), (1, 7))

    def test_path_small_min(self):
        path = tuple(get_triangle_path(StringIO(self.TR_SMALL), lambda x: -x))
        assert path == ((0, 5), (1, 6), (1, 6), (2, 1))

    def test_path_ws_max(self):
        path = tuple(get_triangle_path(StringIO(self.TR_WS)))
        assert path == ((0, 5), (0, 9), (1, 6), (1, 7))

    def test_path_ws_min(self):
        path = tuple(get_triangle_path(StringIO(self.TR_WS), lambda x: -x))
        assert path == ((0, 5), (1, 6), (1, 6), (2, 1))

    def test_path_dupes_max(self):
        path = tuple(get_triangle_path(StringIO(self.TR_DUPES)))
        assert path == ((0, 4), (0, 7), (0, 3), (1, 1))

    def test_path_dupes_min(self):
        path = tuple(get_triangle_path(StringIO(self.TR_DUPES), lambda x: -x))
        assert path == ((0, 4), (1, 6), (1, 3), (1, 1))

    def test_values_small_max(self):
        values = tuple(get_triangle_values(StringIO(self.TR_SMALL)))
        assert values == (5, 9, 6, 7)

    def test_values_small_min(self):
        values = tuple(get_triangle_values(StringIO(self.TR_SMALL), lambda x: -x))
        assert values == (5, 6, 6, 1)


if __name__ == '__main__':
    if len(sys.argv) > 1 and sys.argv[1] == 'test':
        unittest.main(argv=sys.argv[0:1])
    else:
        # by default we print the sum of the input triangle
        max_sum = sum(get_triangle_values(sys.stdin))
        print 'Max sum: %d' % max_sum