day9.py

#!/bin/python

with open('day9.txt') as f:
    heightmap = [list(map(int, line.strip())) for line in f]

def get_neighbours(p):
    neighbours = []
    (x, y) = p[0], p[1]
    if x != 0: neighbours.append((x-1, y))
    if x != len(heightmap[y])-1: neighbours.append((x+1, y))
    if y != 0: neighbours.append((x, y-1))
    if y != len(heightmap)-1: neighbours.append((x, y+1))
    return neighbours

def is_lowpoint(p):
    (x, y) = p[0], p[1]
    return all([heightmap[y][x] < heightmap[n[1]][n[0]] for n in get_neighbours(p)])

lowpoints = [(x, y) for y in range(len(heightmap)) for x in range(len(heightmap[0])) if is_lowpoint((x, y))]
risklevel = sum([heightmap[lp[1]][lp[0]]+1 for lp in lowpoints])
print('Part One:', risklevel)

def get_basin(p, basin=None):
    if basin is None: basin = {}
    (x, y) = p[0], p[1]
    if heightmap[y][x] == 9: return basin
    basin[(x, y)] = 1
    for n in get_neighbours(p):
        if not n in basin:
            get_basin(n, basin)
    return basin

def multiply(values):
    r = 1
    for v in values:
        r *= v
    return r

basins = [sum(get_basin((lp)).values()) for lp in lowpoints]
print('Part Two:', multiply(sorted(basins, reverse=True)[:3]))