commit 0b23461c5d9d68fc0084956f2eb70b31822ca1c1
parent 3e71fb97cb190e8e87d43adbcde6dc8bb342b1d5
Author: Eamon Caddigan <eamon.caddigan@gmail.com>
Date: Thu, 9 Dec 2021 14:05:09 -0500
Solution to day 9, part 2
Diffstat:
A | day09_part2.py | | | 65 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
1 file changed, 65 insertions(+), 0 deletions(-)
diff --git a/day09_part2.py b/day09_part2.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python
+"""Advent of Code 2021, day 9 (part 2): find the three largest "basins" in the
+heightmap"""
+
+# I was wrong about convolution, but I'm still glad I'm working with a 2D
+# array. I'm not convinced that working from the 'low points' to find the size
+# of the basins is the most algorithmically efficient approach, but it's
+# conceptually simple: for each of our low points (226, in my case, and 4 in
+# the case of the example input), we'll perform depth first search to determine
+# which nodes (here, positions in the 2D array) belong to that basin and also
+# find its size. Doing depth first search in "pure" functional languages is
+# acknowledged as a hard problem, so I'm not going to even bother trying in
+# Python--I'll freely take advantage of side effects and not feel even a tiny
+# bit of guilt for it (okay, maybe a tiny bit).
+
+import numpy as np
+from day09_part1 import (EXAMPLE_INPUT,
+ convert_input_to_array,
+ find_low_points)
+from utils import get_puzzle_input
+
+def find_basin_sizes(heightmap):
+ """Return a sorted list of basin sizes. Note that the number of unique
+ basins won't necessarily equal the number of low points (two local minima
+ can easily share a basin), and since we don't need to worry about
+ preserving the mapping between low points and basins we won't even try"""
+ low_points = find_low_points(heightmap)
+ visitmap = np.zeros_like(low_points) # 2D bool array initialized False
+
+ def visit_node_and_get_size(row, col):
+ """Depth first search: mark the `visitmap` visited at the current row
+ and column, and return the size (so far) of the basin. Returns 0 when
+ encountering a basin "ridge" (a location where `heightmap[row, col] ==
+ 9`, a previously visited location (which means that some low points
+ will return sizes of 0 when they are part of previously-explored
+ basins), or a node beyond the edge of `heightmap`."""
+ if (row < 0 or row >= heightmap.shape[0]) or \
+ (col < 0 or col >= heightmap.shape[1]) or \
+ visitmap[row, col] or \
+ heightmap[row, col] == 9:
+ return 0
+ visitmap[row, col] = True
+ return 1 + \
+ visit_node_and_get_size(row-1, col) + \
+ visit_node_and_get_size(row+1, col) + \
+ visit_node_and_get_size(row, col-1) + \
+ visit_node_and_get_size(row, col+1)
+
+ basin_sizes = [visit_node_and_get_size(r, c) for \
+ r, c in zip(*np.where(low_points))]
+
+ return sorted(basin_sizes, reverse=True)
+
+def solve_puzzle(input_string):
+ """Return the numeric solution to the puzzle"""
+ return np.prod(find_basin_sizes(convert_input_to_array(input_string))[0:3])
+
+def main():
+ """Run when the file is called as a script"""
+ assert solve_puzzle(EXAMPLE_INPUT) == 1134
+ print("Product of the sizes of the largest three basins:",
+ solve_puzzle(get_puzzle_input(9)))
+
+if __name__ == "__main__":
+ main()