advent_of_code_2021

day12_part2.py (3599B)

---

 #!/usr/bin/env python
 """Advent of Code 2021, day 12 (part 2): Passage Pathing
 This problem is graph search, with the wrinkle that "large caves" can be
 visited more than once, and, in an update to part 1, each path can visit one
 small cave up to twice."""
 
 # I think I'll only need to make minor changes to my part 1 code, but
 # unfortunately I'll need to copy-paste stuff to use it.
 
 from typing import Tuple, List
 import numpy as np
 from day12_part1 import (EXAMPLE_INPUT,
                          parse_input,
                          update_visits)
 from utils import get_puzzle_input
 
 # This is the only function that changed
 
 def is_visitable(node_names: np.ndarray, node_visits: np.ndarray) -> np.ndarray:
     """Returns a boolean array indicating which nodes are still visitable
     during search"""
     # All 'large' carerns are always visitable. If any small cavern has been
     # visited twice, all small caverns are only visitable if they haven't been
     # visited yet, otherwise, small caverns are visitable if they've been
     # visited 0 or 1 times
     visitable_nodes = np.char.isupper(node_names)
     if (node_visits[~visitable_nodes] > 1).any():
         visitable_nodes |= node_visits == 0
     else:
         visitable_nodes |= node_visits <= 1
     visitable_nodes &= node_names != "start"
     return visitable_nodes
 
 # The definitions of the below functions are the same as part 1, but they need
 # to be redefined here so that they call the right `is_visitable`.
 
 def traverse_node(node_id: int,
                   node_names: np.ndarray,
                   node_visits: np.ndarray,
                   edges: np.ndarray) -> List[Tuple[int, ...]]:
     """Perform depth first search through the given node. Search terminates at
     a node that cannot visit any neighbors or a node that's named 'end'.
     Returns a list of paths, each path is a tuple of node IDs."""
     # End of search when we hit 'end'
     if node_names[node_id] == "end":
         return [(node_id,)]
 
     updated_node_visits = update_visits(node_id, node_visits)
     visitable_nodes = (edges[node_id,] > 0) & is_visitable(node_names, updated_node_visits)
 
     # Also end of search when (if) we hit somebody with no neighbors
     if not visitable_nodes.any():
         return[(node_id,)]
 
     # Continue the search
     paths = []
     for neighbor_id in visitable_nodes.nonzero()[0]:
         for path in traverse_node(neighbor_id, node_names, updated_node_visits, edges):
             paths.append((node_id,) + path)
     return paths
 
 def find_paths_start_to_end(node_names: np.ndarray,
                             edges: np.ndarray) -> List[Tuple[int, ...]]:
     """Return a list of all paths from the node named 'start' to a node named
     'end' (paths are tuples of node IDs)"""
     start_node = int((node_names == 'start').nonzero()[0][0])
     end_node = int((node_names == 'end').nonzero()[0][0])
     return [path for path in \
             traverse_node(start_node,
                           node_names,
                           np.zeros(node_names.size, dtype=int),
                           edges) \
             if path[-1] == end_node]
 
 def solve_puzzle(input_string: str) -> int:
     """Return the numeric solution to the puzzle"""
     return len(find_paths_start_to_end(*parse_input(input_string)))
 
 def main() -> None:
     """Run when the file is called as a script"""
     assert solve_puzzle(EXAMPLE_INPUT) == 3509
     print("Number of routes from 'start' to 'end'",
           "(with up to two visits per small cavern):",
           solve_puzzle(get_puzzle_input(12)))
 
 if __name__ == "__main__":
     main()