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

day12_part2.py (3599B)


      1 #!/usr/bin/env python
      2 """Advent of Code 2021, day 12 (part 2): Passage Pathing
      3 This problem is graph search, with the wrinkle that "large caves" can be
      4 visited more than once, and, in an update to part 1, each path can visit one
      5 small cave up to twice."""
      6 
      7 # I think I'll only need to make minor changes to my part 1 code, but
      8 # unfortunately I'll need to copy-paste stuff to use it.
      9 
     10 from typing import Tuple, List
     11 import numpy as np
     12 from day12_part1 import (EXAMPLE_INPUT,
     13                          parse_input,
     14                          update_visits)
     15 from utils import get_puzzle_input
     16 
     17 # This is the only function that changed
     18 
     19 def is_visitable(node_names: np.ndarray, node_visits: np.ndarray) -> np.ndarray:
     20     """Returns a boolean array indicating which nodes are still visitable
     21     during search"""
     22     # All 'large' carerns are always visitable. If any small cavern has been
     23     # visited twice, all small caverns are only visitable if they haven't been
     24     # visited yet, otherwise, small caverns are visitable if they've been
     25     # visited 0 or 1 times
     26     visitable_nodes = np.char.isupper(node_names)
     27     if (node_visits[~visitable_nodes] > 1).any():
     28         visitable_nodes |= node_visits == 0
     29     else:
     30         visitable_nodes |= node_visits <= 1
     31     visitable_nodes &= node_names != "start"
     32     return visitable_nodes
     33 
     34 # The definitions of the below functions are the same as part 1, but they need
     35 # to be redefined here so that they call the right `is_visitable`.
     36 
     37 def traverse_node(node_id: int,
     38                   node_names: np.ndarray,
     39                   node_visits: np.ndarray,
     40                   edges: np.ndarray) -> List[Tuple[int, ...]]:
     41     """Perform depth first search through the given node. Search terminates at
     42     a node that cannot visit any neighbors or a node that's named 'end'.
     43     Returns a list of paths, each path is a tuple of node IDs."""
     44     # End of search when we hit 'end'
     45     if node_names[node_id] == "end":
     46         return [(node_id,)]
     47 
     48     updated_node_visits = update_visits(node_id, node_visits)
     49     visitable_nodes = (edges[node_id,] > 0) & is_visitable(node_names, updated_node_visits)
     50 
     51     # Also end of search when (if) we hit somebody with no neighbors
     52     if not visitable_nodes.any():
     53         return[(node_id,)]
     54 
     55     # Continue the search
     56     paths = []
     57     for neighbor_id in visitable_nodes.nonzero()[0]:
     58         for path in traverse_node(neighbor_id, node_names, updated_node_visits, edges):
     59             paths.append((node_id,) + path)
     60     return paths
     61 
     62 def find_paths_start_to_end(node_names: np.ndarray,
     63                             edges: np.ndarray) -> List[Tuple[int, ...]]:
     64     """Return a list of all paths from the node named 'start' to a node named
     65     'end' (paths are tuples of node IDs)"""
     66     start_node = int((node_names == 'start').nonzero()[0][0])
     67     end_node = int((node_names == 'end').nonzero()[0][0])
     68     return [path for path in \
     69             traverse_node(start_node,
     70                           node_names,
     71                           np.zeros(node_names.size, dtype=int),
     72                           edges) \
     73             if path[-1] == end_node]
     74 
     75 def solve_puzzle(input_string: str) -> int:
     76     """Return the numeric solution to the puzzle"""
     77     return len(find_paths_start_to_end(*parse_input(input_string)))
     78 
     79 def main() -> None:
     80     """Run when the file is called as a script"""
     81     assert solve_puzzle(EXAMPLE_INPUT) == 3509
     82     print("Number of routes from 'start' to 'end'",
     83           "(with up to two visits per small cavern):",
     84           solve_puzzle(get_puzzle_input(12)))
     85 
     86 if __name__ == "__main__":
     87     main()