Do day 9

[max/advent_of_code_2021.git] / 9_2.py

diff --git a/9_2.py b/9_2.py
new file mode 100644 (file)
index 0000000..31780c2
--- /dev/null
+++ b/9_2.py
@@ -0,0 +1,49 @@
+"""Find where smoke settles"""
+import math
+
+heights = []
+
+with open("9_input.txt") as f:
+    for line in f:
+        heights.append([int(char) for char in line.strip()])
+
+n_rows = len(heights)
+n_cols = len(heights[0])
+
+mins = []
+mins_coords = []
+for i in range(n_rows):
+    for j in range(n_cols):
+        val = heights[i][j]
+        to_compare = [
+            heights[i - 1][j] if i != 0 else 10,
+            heights[i][j - 1] if j != 0 else 10,
+            10 if i == n_rows - 1 else heights[i + 1][j],
+            10 if j == n_cols - 1 else heights[i][j + 1],
+        ]
+        if val < min(to_compare):
+            mins.append(val)
+            mins_coords.append((i, j))
+
+basins = []
+for minimum in mins_coords:
+    to_walk = set([minimum])
+    walked = set()
+
+    while to_walk:
+        i, j = to_walk.pop()
+        orth_coords = [(i - 1, j), (i, j - 1), (i + 1, j), (i, j + 1)]
+        for x, y in orth_coords:
+            try:
+                if (x >= 0 and y >= 0 and heights[x][y] != 9 and
+                        (x, y) not in walked):
+                    to_walk.add((x, y))
+            except IndexError:
+                pass
+
+        walked.add((i, j))
+
+    basins.append(walked)
+
+basins = sorted(basins, key=lambda x: len(x))
+print(math.prod([len(b) for b in basins[-3:]]))