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

day15_part1.py (4115B)


      1 #!/usr/bin/env python
      2 """Advent of Code 2021, day 15 (part 1): Chiton
      3 Another graph search problem, but this one requires a bit more care"""
      4 
      5 # Since I'm not TRYING to reinvent the wheel, I'm going to see if any of
      6 # scipy's graph search algorithms can handle this. I'll implement A* myself if
      7 # I have to, but I'll try to avoid it.
      8 
      9 #from typing import List, Tuple, cast
     10 import numpy as np
     11 from scipy.sparse import csr_matrix
     12 from scipy.sparse.csgraph import dijkstra
     13 from utils import get_puzzle_input, convert_input_to_array
     14 
     15 EXAMPLE_INPUT = \
     16 """1163751742
     17 1381373672
     18 2136511328
     19 3694931569
     20 7463417111
     21 1319128137
     22 1359912421
     23 3125421639
     24 1293138521
     25 2311944581
     26 """
     27 
     28 def grid_to_distance_matrix(input_array: np.ndarray) -> csr_matrix:
     29     """Given a n x m grid of node values, return a sparse array (representing
     30     [n x m] x [n x m] associations) giving the connectivity/distance matrix"""
     31     # Remember that numpy stores data in row-major order, so I'm preserving the
     32     # indicies into the original 2d array in the distance matrix: e.g., the
     33     # last element in a 3x3 array has index 8 (thanks to zero-indexing), and is
     34     # connected to index 7 (to its left) and index 5 (above it), so the sparse
     35     # distance matrix will have connections in row 8 at columns 5 and 7. Note
     36     # that this is a _directed_ distance/connectivity/association matrix,
     37     # because the cost of going from #8 to #7 isn't necessarily the same as
     38     # going from #7 to #8.
     39 
     40     # In the (sparse) distance matrix, data will be stored as:
     41     # Row: node id (row-major indexing), Col: neighbor id
     42     # We'll start off by pretending that every node has four neighbors
     43     dist_row_ids = np.repeat(np.arange(input_array.size, dtype=int), 4)
     44     dist_col_ids = np.empty_like(dist_row_ids)
     45     dist_col_ids[0::4] = dist_row_ids[0::4] - 1                     # Left
     46     dist_col_ids[1::4] = dist_row_ids[1::4] + 1                     # Right
     47     dist_col_ids[2::4] = dist_row_ids[2::4] - input_array.shape[1]  # Top
     48     dist_col_ids[3::4] = dist_row_ids[3::4] + input_array.shape[1]  # Bottom
     49 
     50     # Now we remove the invalid neighbors (at the edges)
     51     grid_col_ids = dist_row_ids % input_array.shape[1]
     52     grid_row_ids = dist_row_ids // input_array.shape[1]
     53     keep_neighbor = np.empty_like(grid_col_ids, dtype=bool)
     54     keep_neighbor[0::4] = grid_col_ids[0::4] >= 1
     55     keep_neighbor[1::4] = grid_col_ids[1::4] < (input_array.shape[1] - 1)
     56     keep_neighbor[2::4] = grid_row_ids[2::4] >= 1
     57     keep_neighbor[3::4] = grid_row_ids[3::4] < (input_array.shape[0] - 1)
     58     dist_row_ids = dist_row_ids[keep_neighbor]
     59     dist_col_ids = dist_col_ids[keep_neighbor]
     60 
     61     # There are a bunch of ways to instantiate a csr_matrix, this is the most
     62     # appropriate for this data format
     63     return csr_matrix((np.reshape(input_array, input_array.size)[dist_col_ids],
     64                        (dist_row_ids, dist_col_ids)))
     65 
     66 def find_shortest_path(distance_matrix: csr_matrix):
     67     """Runs scipy's implementation of Dijkstra's algorithm (which isn't as fast
     68     as A* in theory, but since this is optimized and my hacked together A*
     69     wouldn't be, probably is faster in practice) and returns the distance from
     70     the starting point to all other points, and the index of the preceding
     71     point"""
     72     return dijkstra(distance_matrix,
     73                     directed=True,              # This is a directed graph
     74                     indices=0,                  # Find the distance from here
     75                     return_predecessors=True,   # Reconstructable path
     76                     min_only=True)[0:2]         # Just want this distance
     77 
     78 def solve_puzzle(input_string: str) -> int:
     79     """Return the numeric solution to the puzzle"""
     80     return int(
     81         find_shortest_path(
     82             grid_to_distance_matrix(
     83                 convert_input_to_array(input_string)
     84             )
     85         )[0][-1]
     86     )
     87 
     88 def main() -> None:
     89     """Run when the file is called as a script"""
     90     assert solve_puzzle(EXAMPLE_INPUT) == 40
     91     print("Total risk of best path:",
     92           solve_puzzle(get_puzzle_input(15)))
     93 
     94 if __name__ == "__main__":
     95     main()