diff options
Diffstat (limited to '9.py')
| -rw-r--r-- | 9.py | 44 |
1 files changed, 21 insertions, 23 deletions
@@ -1,45 +1,43 @@ from functools import reduce -from utils import open_day +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(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 +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(x, y): +def basin_size(p): seen = set() - queue = deque(((x, y),)) + queue = deque((p,)) size = 0 while queue: - x, y = queue.popleft() - if (x, y) in seen: + p = queue.popleft() + if p 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: + seen.add(p) + if heatmap[p.y][p.x] == 9: continue size += 1 - for dx, dy in adj: - queue.append((x + dx, y + dy)) + for n in adjacent_bounded(p, bounds, False): + # print(p, n) + queue.append(n) return size totrisk = 0 lows = [] -for y in range(len(heatmap)): - for x in range(len(heatmap[y])): - if is_low(x, y): +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(x, y)) + lows.append(basin_size(p)) print(totrisk) lows.sort() print(reduce(lambda x, y: x * y, lows[-3:], 1)) |