day12_part1.py (4672B)
1 #!/usr/bin/env python 2 """Advent of Code 2021, day 12 (part 1): Passage Pathing 3 This problem is graph search, with the wrinkle that "large caves" can be 4 visited more than once""" 5 6 # We'll find all the paths from `start` to `end` in order to count them. This 7 # is basically a very straightforward graph traversal, with the caveat that 8 # "large caves" can be visited more than once--I'm not going to worry about 9 # getting caught in loops between them (and if I have to change that in code, I 10 # hope I remember to edit this comment). Since this is my year of pandas and 11 # numpy, I am going to implement edges with a 2d numpy array (even though the 12 # edge weights are simple--1 or 0--and symmetrical), and node information with 13 # 1d numpy arrays. 14 15 from typing import Tuple, List 16 import numpy as np 17 from utils import get_puzzle_input 18 19 EXAMPLE_INPUT = \ 20 """fs-end 21 he-DX 22 fs-he 23 start-DX 24 pj-DX 25 end-zg 26 zg-sl 27 zg-pj 28 pj-he 29 RW-he 30 fs-DX 31 pj-RW 32 zg-RW 33 start-pj 34 he-WI 35 zg-he 36 pj-fs 37 start-RW 38 """ 39 40 def parse_input(input_string: str) -> Tuple[np.ndarray, np.ndarray]: 41 """Read the cave map, and return a tuple containing two arrays: 42 * A 1d array of node names (strings) 43 * A 2d array (number of nodes x number of nodes) to represent edges""" 44 # Read in the data and find the nodes and edges 45 node_set = set() 46 edge_list = [] 47 for edge_string in input_string.rstrip('\n').split('\n'): 48 edge_nodes_set = set(edge_string.split('-')) 49 node_set.update(edge_nodes_set) 50 edge_list.append(edge_nodes_set) 51 52 # Create the numpy objects 53 node_names = np.array(list(node_set)) 54 edges = np.zeros((node_names.size, node_names.size), dtype=int) 55 56 # Fill in the edges 57 # (Here's where we're punished for representing the data this way, other 58 # models would be much easier to parse) 59 for edge_nodes_set in edge_list: 60 n_1, n_2 = edge_nodes_set 61 edges[node_names == n_1, node_names == n_2] = 1 62 edges[node_names == n_2, node_names == n_1] = 1 63 64 return (node_names, edges) 65 66 def is_visitable(node_names: np.ndarray, node_visits: np.ndarray) -> np.ndarray: 67 """Returns a boolean array indicating which nodes are still visitable 68 during search (defined as nodes associated with 'large' caverns or nodes 69 not yet visited)""" 70 return (node_visits == 0) | np.char.isupper(node_names) 71 72 def update_visits(node_id: int, node_visits: np.ndarray) -> np.ndarray: 73 """Return a copy of the visit counts with the specified node marked as 74 visited again""" 75 node_visits_copy = node_visits.copy() 76 node_visits_copy[node_id] += 1 77 return node_visits_copy 78 79 def traverse_node(node_id: int, 80 node_names: np.ndarray, 81 node_visits: np.ndarray, 82 edges: np.ndarray) -> List[Tuple[int, ...]]: 83 """Perform depth first search through the given node. Search terminates at 84 a node that cannot visit any neighbors or a node that's named 'end'. 85 Returns a list of paths, each path is a tuple of node IDs.""" 86 # End of search when we hit 'end' 87 if node_names[node_id] == "end": 88 return [(node_id,)] 89 90 updated_node_visits = update_visits(node_id, node_visits) 91 visitable_nodes = (edges[node_id,] > 0) & is_visitable(node_names, updated_node_visits) 92 93 # Also end of search when (if) we hit somebody with no neighbors 94 if not visitable_nodes.any(): 95 return[(node_id,)] 96 97 # Continue the search 98 paths = [] 99 for neighbor_id in visitable_nodes.nonzero()[0]: 100 for path in traverse_node(neighbor_id, node_names, updated_node_visits, edges): 101 paths.append((node_id,) + path) 102 return paths 103 104 def find_paths_start_to_end(node_names: np.ndarray, 105 edges: np.ndarray) -> List[Tuple[int, ...]]: 106 """Return a list of all paths from the node named 'start' to a node named 107 'end' (paths are tuples of node IDs)""" 108 start_node = int((node_names == 'start').nonzero()[0][0]) 109 end_node = int((node_names == 'end').nonzero()[0][0]) 110 return [path for path in \ 111 traverse_node(start_node, 112 node_names, 113 np.zeros(node_names.size, dtype=int), 114 edges) \ 115 if path[-1] == end_node] 116 117 def solve_puzzle(input_string: str) -> int: 118 """Return the numeric solution to the puzzle""" 119 return len(find_paths_start_to_end(*parse_input(input_string))) 120 121 def main() -> None: 122 """Run when the file is called as a script""" 123 assert solve_puzzle(EXAMPLE_INPUT) == 226 124 print("Number of routes from 'start' to 'end':", 125 solve_puzzle(get_puzzle_input(12))) 126 127 if __name__ == "__main__": 128 main()