day15_part2.py (1733B)
1 #!/usr/bin/env python 2 """Advent of Code 2021, day 15 (part 2): Chiton 3 Another graph search problem, but this one requires even more more care""" 4 5 # Uh oh, is scipy's implementation of Dijkstra's algorithm fast enough to do a 6 # humongous graph? I certainly hope so! 7 8 #from typing import List, Tuple, cast 9 import numpy as np 10 from day15_part1 import (EXAMPLE_INPUT, 11 grid_to_distance_matrix, 12 find_shortest_path) 13 from utils import get_puzzle_input, convert_input_to_array 14 15 def quintuple_grid(input_array: np.ndarray) -> np.ndarray: 16 """Quintuple the input grid in each direction, incrementing the values and 17 using modular arithmetic to stay in the range [1, 9]""" 18 output_array_col = input_array.copy() 19 for row in range(1, 5): 20 output_array_col = np.concatenate((output_array_col, input_array + row), 21 axis=0) 22 output_array = output_array_col.copy() 23 for col in range(1, 5): 24 output_array = np.concatenate((output_array, output_array_col + col), 25 axis=1) 26 return ((output_array - 1) % 9) + 1 27 28 def solve_puzzle(input_string: str) -> int: 29 """Return the numeric solution to the puzzle""" 30 return int( 31 find_shortest_path( 32 grid_to_distance_matrix( 33 quintuple_grid( 34 convert_input_to_array(input_string) 35 ) 36 ) 37 )[0][-1] 38 ) 39 40 def main() -> None: 41 """Run when the file is called as a script""" 42 assert solve_puzzle(EXAMPLE_INPUT) == 315 43 print("Total risk of best path (using larger map):", 44 solve_puzzle(get_puzzle_input(15))) 45 46 if __name__ == "__main__": 47 main()