advent_of_code_2021

day09_part2.py (3043B)

---

 #!/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()