day9.py

from utils import Pt, Grd
import math

grd = Grd()
for i, l in enumerate(LINES):
    for j, c in enumerate(l):
        grd[Pt(i, j)] = int(c)


def expand(pt):
    basin = {pt}
    queue = [pt]

    while queue:
        p = queue.pop()
        for nb in p.nb4:
            if nb not in basin and nb in grd and grd[nb] < 9 and grd[nb] > grd[p]:
                basin.add(nb)
                queue.append(nb)

    return basin


tot = 0
basins = []
for pt in grd:
    if all(n not in grd or grd[pt] < grd[n] for n in pt.nb4):
        tot += grd[pt] + 1
        basins.append(len(expand(pt)))

print("part 1", tot)
print("part 2", math.prod(list(sorted(basins))[-3:]))