from functools import reduce from utils import open_day, Point2D, adjacent_bounded from collections import deque heatmap = [[int(c) for c in l.rstrip()] for l in open_day(9)] bounds = Point2D(len(heatmap[0]), len(heatmap)) adj = [(1, 0), (0, 1), (-1, 0), (0, -1)] def is_low(p): risk = heatmap[p.y][p.x] for n in adjacent_bounded(p, bounds, False): if heatmap[n.y][n.x] <= risk: return False return True def basin_size(p): seen = set() queue = deque((p,)) size = 0 while queue: p = queue.popleft() if p in seen: continue seen.add(p) if heatmap[p.y][p.x] == 9: continue size += 1 for n in adjacent_bounded(p, bounds, False): # print(p, n) queue.append(n) return size totrisk = 0 lows = [] for y in range(bounds.y): for x in range(bounds.x): p = Point2D(x, y) if is_low(p): totrisk += 1 + heatmap[y][x] lows.append(basin_size(p)) print(totrisk) lows.sort() print(reduce(lambda x, y: x * y, lows[-3:], 1))