1 files changed, 45 insertions, 0 deletions
diff --git a/9.py b/9.py
new file mode 100644
index 0000000..d1a63b9
--- /dev/null
+++ b/9.py
@@ -0,0 +1,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))