about summary refs log tree commit diff
path: root/aoc/2021/09/part-two.py
blob: 99524c499e9d82ea1f0a6780c635708fe7c215dd (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from collections import deque
from functools import reduce
from operator import mul
from sys import stdin

heightmap = [[int(c) for c in s] for s in stdin.read().strip().split()]
get = lambda y, x: 10 if min(x, y) < 0 or max(x, y) > 99 else heightmap[y][x]
adjacent = lambda y, x: ((y, x + 1), (y + 1, x), (y, x - 1), (y - 1, x))
basins = {(i, j): 0 for i, row in enumerate(heightmap)
          for j, height in enumerate(row)
          if all(height < get(y, x) for y, x in adjacent(i, j))}
queue, visited = deque((key, key) for key in basins), set()
while queue:
    key, pos = queue.popleft()
    if pos in visited: continue
    basins[key] += 1
    visited.add(pos)
    for adj in adjacent(*pos):
        if get(*adj) < 9 and adj not in visited: queue.append((key, adj))
print(reduce(mul, sorted(basins.values(), reverse=True)[:3]))