AoC 2021 Day 9

import os
from collections import defaultdict


def read_and_init(fn):
    with open(fn, 'r') as f:
        lines = f.read().splitlines()
    height_map = defaultdict(int)
    coord = 0j
    for line in lines:
        for c in line:
            height_map[coord] = int(c)
            coord += 1
        coord = 0+coord.imag*1j
        coord += 1j
    return height_map


def product(vals):
    result = 1
    for v in vals:
        result *= v
    return result


def in_bounds(loc: complex, lower: complex, upper: complex) -> bool:
    return (lower.real <= loc.real < upper.real) and (lower.imag <= loc.imag < upper.imag)


def check_local_minima(hmap: defaultdict, loc: complex) -> bool:
    dirs = (-1j, -1, 1, 1j)
    upper = int(max(key.real for key in hmap) + 1)+1j*int(max(key.imag for key in hmap) + 1)
    return all(hmap[loc + d] > hmap[loc] for d in dirs if in_bounds(loc + d, 0j, upper))


def find_basins(hmap):
    upper = int(max(key.real for key in hmap) + 1) + 1j * int(max(key.imag for key in hmap) + 1)
    minima = get_minima(hmap)
    basins = {}
    dirs = (-1j, -1, 1, 1j)
    for m in minima:
        basins[m] = set()
        queue = [m]
        while queue:
            current = queue.pop(0)
            for neighbor in [current + d for d in dirs if in_bounds(current + d, 0j, upper)]:
                if hmap[neighbor] == 9:
                    continue
                if neighbor not in basins[m]:
                    basins[m].add(neighbor)
                    queue.append(neighbor)
    return basins


def get_minima(hmap):
    upper = int(max(key.real for key in hmap) + 1) + 1j * int(max(key.imag for key in hmap) + 1)
    minima = []
    for y in range(int(upper.imag)):
        for x in range(int(upper.real)):
            if check_local_minima(hmap, x + 1j * y):
                minima.append(x + 1j * y)
    return minima


def part1(hmap):
    minima = get_minima(hmap)
    return sum(hmap[m] for m in minima)


def part2(hmap):
    basins = find_basins(hmap)
    basin_sizes = [len(basins[m]) for m in basins]
    basin_sizes.sort(reverse=True)
    return product(basin_sizes[:3])


h = read_and_init(os.path.dirname(__file__) + '\\in.txt')
res = part1(h)
print('Part 1:', res)
res = part2(h)
print('Part 2:', res)