12_2.py

   1 """ Count number of paths through caves """
   2 import pprint
   3 from collections import defaultdict
   4
   5 def walk_from_node(current_node, path_taken_so_far, visited_small_cave_twice=False):
   6     possible_paths = []
   7     for node in caves[current_node]:
   8         if node == "start":
   9             pass
  10         elif node == "end":
  11             possible_paths.append(path_taken_so_far + ["end"])
  12         elif node == node.lower() and node in path_taken_so_far:
  13             if not visited_small_cave_twice:
  14                 possible_paths += walk_from_node(node, path_taken_so_far + [current_node], True)
  15         else:
  16             possible_paths += walk_from_node(node, path_taken_so_far + [current_node], visited_small_cave_twice)
  17
  18     return possible_paths
  19
  20 caves = defaultdict(set)
  21
  22 with open("12_input.txt") as f:
  23     for line in f:
  24         node_a, node_b = line.strip().split("-")
  25
  26         caves[node_a].add(node_b)
  27         caves[node_b].add(node_a)
  28
  29 walked_paths = walk_from_node("start", [])
  30 print(len(walked_paths))