advent_of_code_2021

day15_part1.py (4115B)

---

 #!/usr/bin/env python
 """Advent of Code 2021, day 15 (part 1): Chiton
 Another graph search problem, but this one requires a bit more care"""
 
 # Since I'm not TRYING to reinvent the wheel, I'm going to see if any of
 # scipy's graph search algorithms can handle this. I'll implement A* myself if
 # I have to, but I'll try to avoid it.
 
 #from typing import List, Tuple, cast
 import numpy as np
 from scipy.sparse import csr_matrix
 from scipy.sparse.csgraph import dijkstra
 from utils import get_puzzle_input, convert_input_to_array
 
 EXAMPLE_INPUT = \
 """1163751742
 1381373672
 2136511328
 3694931569
 7463417111
 1319128137
 1359912421
 3125421639
 1293138521
 2311944581
 """
 
 def grid_to_distance_matrix(input_array: np.ndarray) -> csr_matrix:
     """Given a n x m grid of node values, return a sparse array (representing
     [n x m] x [n x m] associations) giving the connectivity/distance matrix"""
     # Remember that numpy stores data in row-major order, so I'm preserving the
     # indicies into the original 2d array in the distance matrix: e.g., the
     # last element in a 3x3 array has index 8 (thanks to zero-indexing), and is
     # connected to index 7 (to its left) and index 5 (above it), so the sparse
     # distance matrix will have connections in row 8 at columns 5 and 7. Note
     # that this is a _directed_ distance/connectivity/association matrix,
     # because the cost of going from #8 to #7 isn't necessarily the same as
     # going from #7 to #8.
 
     # In the (sparse) distance matrix, data will be stored as:
     # Row: node id (row-major indexing), Col: neighbor id
     # We'll start off by pretending that every node has four neighbors
     dist_row_ids = np.repeat(np.arange(input_array.size, dtype=int), 4)
     dist_col_ids = np.empty_like(dist_row_ids)
     dist_col_ids[0::4] = dist_row_ids[0::4] - 1                     # Left
     dist_col_ids[1::4] = dist_row_ids[1::4] + 1                     # Right
     dist_col_ids[2::4] = dist_row_ids[2::4] - input_array.shape[1]  # Top
     dist_col_ids[3::4] = dist_row_ids[3::4] + input_array.shape[1]  # Bottom
 
     # Now we remove the invalid neighbors (at the edges)
     grid_col_ids = dist_row_ids % input_array.shape[1]
     grid_row_ids = dist_row_ids // input_array.shape[1]
     keep_neighbor = np.empty_like(grid_col_ids, dtype=bool)
     keep_neighbor[0::4] = grid_col_ids[0::4] >= 1
     keep_neighbor[1::4] = grid_col_ids[1::4] < (input_array.shape[1] - 1)
     keep_neighbor[2::4] = grid_row_ids[2::4] >= 1
     keep_neighbor[3::4] = grid_row_ids[3::4] < (input_array.shape[0] - 1)
     dist_row_ids = dist_row_ids[keep_neighbor]
     dist_col_ids = dist_col_ids[keep_neighbor]
 
     # There are a bunch of ways to instantiate a csr_matrix, this is the most
     # appropriate for this data format
     return csr_matrix((np.reshape(input_array, input_array.size)[dist_col_ids],
                        (dist_row_ids, dist_col_ids)))
 
 def find_shortest_path(distance_matrix: csr_matrix):
     """Runs scipy's implementation of Dijkstra's algorithm (which isn't as fast
     as A* in theory, but since this is optimized and my hacked together A*
     wouldn't be, probably is faster in practice) and returns the distance from
     the starting point to all other points, and the index of the preceding
     point"""
     return dijkstra(distance_matrix,
                     directed=True,              # This is a directed graph
                     indices=0,                  # Find the distance from here
                     return_predecessors=True,   # Reconstructable path
                     min_only=True)[0:2]         # Just want this distance
 
 def solve_puzzle(input_string: str) -> int:
     """Return the numeric solution to the puzzle"""
     return int(
         find_shortest_path(
             grid_to_distance_matrix(
                 convert_input_to_array(input_string)
             )
         )[0][-1]
     )
 
 def main() -> None:
     """Run when the file is called as a script"""
     assert solve_puzzle(EXAMPLE_INPUT) == 40
     print("Total risk of best path:",
           solve_puzzle(get_puzzle_input(15)))
 
 if __name__ == "__main__":
     main()