from lib import get_data, Grid2D from collections import deque data = get_data(__file__).strip() g = Grid2D(data) s = 0 low_points = [] for row in range(g.n_rows): for col in range(g.n_cols): h = int(g[(row, col)]) for nb in g.neighbors_ort((row, col)): if int(g[nb]) <= h: break else: low_points.append((row, col)) s += 1 + h print(s) basins = [] for row, col in low_points: start = None visited = set() queue = deque([(row, col)]) while queue: vertex = queue.popleft() if vertex in visited: continue visited.add(vertex) neighbors = [] for neighbor in g.neighbors_ort(vertex): if neighbor in visited: continue neighbor_height = int(g[neighbor]) if neighbor_height < 9: queue.append(neighbor) basins.append(len(visited)) s = 1 for x in sorted(basins)[-3:]: s *= x print(s)