Solve 2021 day 9

farbfetzen · farbfetzen · commit 1a02d356f6f4 · 2021-12-11T16:02:34.000+01:00
diff --git a/README.md b/README.md
@@ -3,7 +3,7 @@
 
 | Year | Stars |
 |------|------:|
-| 2021 |  16 ⭐ |
+| 2021 |  18 ⭐ |
 | 2020 |  50 ⭐ |
 | 2019 |  44 ⭐ |
 | 2018 |   8 ⭐ |
diff --git a/input/2021-09-sample.txt b/input/2021-09-sample.txt
@@ -0,0 +1,5 @@
+2199943210
+3987894921
+9856789892
+8767896789
+9899965678
diff --git a/python/2021/day09.py b/python/2021/day09.py
@@ -0,0 +1,79 @@
+# https://adventofcode.com/2021/day/9
+
+
+from collections import namedtuple
+from math import prod
+
+
+SAMPLE_PATH = "../../input/2021-09-sample.txt"
+INPUT_PATH = "../../input/2021-09-input.txt"
+Point = namedtuple("Point", ("x", "y"))
+
+
+def get_data(filename):
+    with open(filename) as file:
+        data = file.read().splitlines()
+    heightmap = tuple(tuple(int(x) for x in list(line)) for line in data)
+    low_points = get_low_points(heightmap)
+    return heightmap, low_points
+
+
+def get_neighbors(x, y, max_x, max_y):
+    neighbors = []
+    if x > 0:
+        neighbors.append(Point(x - 1, y))
+    if x < max_x:
+        neighbors.append(Point(x + 1, y))
+    if y > 0:
+        neighbors.append(Point(x, y - 1))
+    if y < max_y:
+        neighbors.append(Point(x, y + 1))
+    return neighbors
+
+
+def get_low_points(heightmap):
+    low_points = []
+    max_x = len(heightmap[0]) - 1
+    max_y = len(heightmap) - 1
+    for y, row in enumerate(heightmap):
+        for x, height in enumerate(row):
+            neighbors = get_neighbors(x, y, max_x, max_y)
+            if all(height < heightmap[n.y][n.x] for n in neighbors):
+                low_points.append(Point(x, y))
+    return low_points
+
+
+def flood_fill(heightmap, origin, max_x, max_y):
+    neighbors = [origin]
+    basin_members = {origin}
+    while neighbors:
+        current = neighbors.pop()
+        new_neighbors = get_neighbors(current.x, current.y, max_x, max_y)
+        for nn in new_neighbors:
+            # Assuming all basins are surrounded by nines, which is true for the sample.
+            if nn not in basin_members and heightmap[nn.y][nn.x] < 9:
+                neighbors.append(nn)
+                basin_members.add(nn)
+    return len(basin_members)
+
+
+def part_1(heightmap, low_points):
+    return sum(heightmap[lp.y][lp.x] + 1 for lp in low_points)
+
+
+def part_2(heightmap, low_points):
+    max_x = len(heightmap[0]) - 1
+    max_y = len(heightmap) - 1
+    basin_sizes = [flood_fill(heightmap, lp, max_x, max_y) for lp in low_points]
+    basin_sizes.sort(reverse=True)
+    return prod(basin_sizes[:3])
+
+
+if __name__ == "__main__":
+    sample_data = get_data(SAMPLE_PATH)
+    assert part_1(*sample_data) == 15
+    assert part_2(*sample_data) == 1134
+
+    challenge_data = get_data(INPUT_PATH)
+    print(part_1(*challenge_data))  # 554
+    print(part_2(*challenge_data))  # 1017792
diff --git a/python/2021/test2021.py b/python/2021/test2021.py
@@ -8,6 +8,7 @@
 import day06
 import day07
 import day08
+import day09
 
 
 class Test2021(unittest.TestCase):
@@ -90,6 +91,15 @@ def test_08(self):
         self.assertEqual(day08.part_1(challenge_data), 352)
         self.assertEqual(day08.part_2(challenge_data), 936117)
 
+    def test_09(self):
+        sample_data = day09.get_data(day09.SAMPLE_PATH)
+        self.assertEqual(day09.part_1(*sample_data), 15)
+        self.assertEqual(day09.part_2(*sample_data), 1134)
+
+        challenge_data = day09.get_data(day09.INPUT_PATH)
+        self.assertEqual(day09.part_1(*challenge_data), 554)
+        self.assertEqual(day09.part_2(*challenge_data), 1017792)
+
 
 if __name__ == "__main__":
     unittest.main()
