advent_12_part2.py

from collections import defaultdict

def is_small_cave(node):
  if node == node.lower():
    return True


def count_paths(node, paths, visited, allow_double_visit):
  if node in visited:
    # If we haven't used our double visit yet, let's use it on this branch.
    # Since otherwise we would need to stop here.
    if node != 'start' and allow_double_visit:
      allow_double_visit = False
    else:
      return 0
  if node == 'end':
    return 1
  if is_small_cave(node):
    visited.add(node)
  valid_paths = 0
  for child in paths[node]:
    # We copy visited, since we're interested on all paths.
    # So each individual branch needs to keep track of only visited nodes
    # in that path.
    valid_paths += count_paths(child, paths, set(visited), allow_double_visit)
  return valid_paths


if __name__ == '__main__':
  paths = defaultdict(list)
  with open('input.txt') as data:
    for line in data:
      path_a, path_b = line.strip().split('-')
      paths[path_a].append(path_b)
      paths[path_b].append(path_a)
    visited = set()
    print(count_paths('start', paths, visited, True))