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