advent/y2021/p09.py

from toolkit import read_image


ans1 = 0
grid = read_image(text)[0]
for p, v in grid.items():
    if all(grid[p + s] > v for s in [1, -1, 1j, -1j] if p + s in grid):
        ans1 += int(v) + 1


valid = {k for k, v in grid.items() if v != '9'}
pending = valid.copy()
basin_sizes = []
while pending:
    edge = {pending.pop()}
    seen = edge.copy()
    while edge:
        edge = {p + s for p in edge for s in [1, -1, 1j, -1j]} & valid - seen
        seen |= edge
    pending -= seen
    basin_sizes.append(len(seen))
*_, x, y, z = sorted(basin_sizes)
ans2 = x * y * z