day9.py

#!/bin/python
# https://adventofcode.com/2020/day/9

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

H, W = len(heightmap), len(heightmap[0])
COORDS = [(x, y) for y in range(H) for x in range(W)]
value_at = lambda p: heightmap[p[0]][p[1]]

def get_neighbours(p):
    x, y = p
    if x > 0:   yield x-1, y
    if x < H-1: yield x+1, y
    if y > 0:   yield x, y-1
    if y < W-1: yield x, y+1
    
def is_lowpoint(p):
    return all(value_at(p) < value_at(n) for n in get_neighbours(p))

lowpoints = list(filter(is_lowpoint, COORDS))
risklevel = sum(value_at(lp) + 1 for lp in lowpoints)
print('Part One:', risklevel)

def get_basin(p, basin=None):
    if basin is None: basin = {}
    if value_at(p) == 9: return basin
    basin[tuple(p)] = 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]))