aoc-2021/9.py

from util import get_input
from functools import reduce

input = get_input("9.input")
input = [[int(a) for a in line] for line in input]

def neighbours(map, x, y):
    ns = []
    if x > 0:
        ns.append(map[x - 1][y])
    if y > 0:
        ns.append(map[x][y - 1])
    if x < len(map) - 1:
        ns.append(map[x + 1][y])
    if y < len(map[0]) - 1:
        ns.append(map[x][y + 1])
    return ns

lows = []

for x in range(len(input)):
    for y in range(len(input[0])):
        ns = neighbours(input, x, y)
        if len([1 for n in ns if n > input[x][y]]) == len(ns):
            lows.append(input[x][y])

print(sum([1 + n for n in lows]))


def basin_size(map, x, y, taken):
    if map[x][y] == 9:
        return 0
    if (x, y) in taken:
        return 0
    size = 1
    taken.add((x, y))
    if x > 0:
        size += basin_size(map, x - 1, y, taken)
    if y > 0:
        size += basin_size(map, x, y - 1, taken)
    if x < len(map) - 1:
        size += basin_size(map, x + 1, y, taken)
    if y < len(map[0]) - 1:
        size += basin_size(map, x, y + 1, taken)
    return size

basins = []
taken = set()

for x in range(len(input)):
    for y in range(len(input[0])):
        size = basin_size(input, x, y, taken)
        if size != 0:
            basins.append(size)

basins.sort()
print(reduce(lambda a, b: a * b, basins[-3:]))