46 lines
1.1 KiB
Python
46 lines
1.1 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
import fileinput
|
||
|
|
import heapq
|
||
|
|
from typing import Generator, Tuple
|
||
|
|
|
||
|
|
heightmap = []
|
||
|
|
for line in fileinput.input():
|
||
|
|
heightmap.append([int(c) for c in line.strip()])
|
||
|
|
|
||
|
|
def neighborhood(x: int, y: int) -> Generator[Tuple[int, int], None, None]:
|
||
|
|
global heightmap
|
||
|
|
if x > 0:
|
||
|
|
yield (x - 1, y)
|
||
|
|
if x < len(heightmap[0]) - 1:
|
||
|
|
yield (x + 1, y)
|
||
|
|
if y > 0:
|
||
|
|
yield (x, y - 1)
|
||
|
|
if y < len(heightmap) - 1:
|
||
|
|
yield (x, y + 1)
|
||
|
|
|
||
|
|
lowpoints = list()
|
||
|
|
for y, row in enumerate(heightmap):
|
||
|
|
for x, height in enumerate(row):
|
||
|
|
if height < min(heightmap[y_][x_] for x_, y_ in neighborhood(x, y)):
|
||
|
|
lowpoints.append((x, y))
|
||
|
|
|
||
|
|
print("part 1:", sum(1 + heightmap[y][x] for x, y in lowpoints))
|
||
|
|
|
||
|
|
basins = []
|
||
|
|
for lowpoint in lowpoints:
|
||
|
|
stack = [lowpoint]
|
||
|
|
visited = set()
|
||
|
|
while stack:
|
||
|
|
x, y = stack.pop()
|
||
|
|
if (x, y) in visited or heightmap[y][x] == 9:
|
||
|
|
continue
|
||
|
|
visited.add((x, y))
|
||
|
|
|
||
|
|
for neighbor in neighborhood(x, y):
|
||
|
|
stack.append(neighbor)
|
||
|
|
|
||
|
|
heapq.heappush(basins, len(visited))
|
||
|
|
|
||
|
|
a, b, c, *_ = heapq.nlargest(3, basins)
|
||
|
|
print("part 2:", a * b * c)
|