advent-of-code/2021/day9/clean.py

62 lines
1.6 KiB
Python
Executable file

#!/usr/bin/env python
from collections import Counter
from dataclasses import dataclass
import math
import sys
from typing import Optional, Tuple
@dataclass
class Point:
height: int
low_point: Optional[Tuple[int, int]]
class Grid:
def __init__(self, data):
self.data = data
self.rows = len(data)
self.cols = len(data[0])
def get(self, x, y):
if 0 <= x < self.rows and 0 <= y < self.cols:
return self.data[x][y]
else:
return None
def find_low_point(self, x, y):
current = self.get(x, y)
if current is None:
raise ValueError
if current.low_point is not None:
return current.low_point
if current.height == 9:
return None
for dx, dy in ((-1, 0), (1, 0), (0, -1), (0, 1)):
target = self.get(x+dx, y+dy)
if target is not None:
if target.height < current.height:
current.low_point = self.find_low_point(x+dx, y+dy)
return current.low_point
return (x, y)
if __name__ == '__main__':
grid = Grid([
[
Point(height=int(n), low_point=None)
for n in i.strip()
]
for i in sys.stdin.readlines()
])
counts = Counter()
for x in range(grid.rows):
for y in range(grid.cols):
low_point = grid.find_low_point(x, y)
if low_point is not None:
counts.update([low_point])
print(sum(grid.get(x, y).height + 1 for x, y in counts.keys()))
print(math.prod(c for _, c in counts.most_common(3)))