advent_of_code_2021

My attempts to work through the 2021 Advent of Code problems.
git clone https://git.eamoncaddigan.net/advent_of_code_2021.git
Log | Files | Refs | README | LICENSE

day09_part1.py (1710B)

      1 #!/usr/bin/env python
      2 """Advent of Code 2021, day 9 (part 1): find the low points of a 2D heightmap"""
      3 
      4 # Ok so I've been really twisting my problems into things that can be solved
      5 # with pandas, but this one really makes the most sense in numpy... and I
      6 # haven't actually used numpy (by itself) that much this AoC. I also suspect
      7 # that part 2 might involve convolutions, but I've been wrong before.
      8 
      9 import numpy as np
     10 from utils import get_puzzle_input
     11 
     12 EXAMPLE_INPUT = \
     13 """2199943210
     14 3987894921
     15 9856789892
     16 8767896789
     17 9899965678
     18 """
     19 
     20 def convert_input_to_array(input_string):
     21     """Convert the input into an array"""
     22     return (
     23         np.array([list(line) for line in input_string.rstrip('\n').split('\n')])
     24         .astype(int)
     25     )
     26 
     27 def find_low_points(heightmap):
     28     """Return a 2d boolean array of local minima"""
     29     return (np.diff(heightmap, axis=0, prepend=10) < 0) & \
     30         np.flip(np.diff(np.flip(heightmap, axis=0), axis=0, prepend=10) < 0,
     31                 axis=0) & \
     32         (np.diff(heightmap, axis=1, prepend=10) < 0) & \
     33         np.flip(np.diff(np.flip(heightmap, axis=1), axis=1, prepend=10) < 0,
     34                 axis=1)
     35 
     36 def calculate_risk_level(heightmap):
     37     """Given a heighmap, find the low points and calculate the risk level"""
     38     return sum(heightmap[find_low_points(heightmap)] + 1)
     39 
     40 def solve_puzzle(input_string):
     41     """Return the numeric solution to the puzzle"""
     42     return calculate_risk_level(convert_input_to_array(input_string))
     43 
     44 def main():
     45     """Run when the file is called as a script"""
     46     assert solve_puzzle(EXAMPLE_INPUT) == 15
     47     print("Risk level of heightmap:",
     48           solve_puzzle(get_puzzle_input(9)))
     49 
     50 if __name__ == "__main__":
     51     main()