1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
|
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))
|
