#!/usr/bin/env python3 import fileinput from timeit import timeit from typing import FrozenSet, List, Optional, Tuple, Union from collections import defaultdict lines = list(fileinput.input()) def main(): # parse graph = defaultdict(list) for line in lines: fro, to = str(line).strip().split("-") graph[fro].append(to) graph[to].append(fro) # part 1 stack: List[Tuple[str, FrozenSet[str]]] = [("start", frozenset())] part1 = 0 while stack: pos, visited = stack.pop() if pos.islower() and pos in visited: continue visited = visited | {pos} if pos == "end": part1 += 1 continue for next_pos in graph[pos]: stack.append((next_pos, visited)) stack2: List[Tuple[str, FrozenSet[str], Optional[str]]] = [ ("start", frozenset(), None) ] part2 = 0 while stack2: pos, visited, doubled = stack2.pop() if pos.islower() and pos in visited: if doubled is None: stack2.append((pos, visited, pos)) continue if doubled != pos: continue doubled = "" visited = visited | {pos} if pos == "end": part2 += 1 continue for next_pos in graph[pos]: if next_pos != "start": stack2.append((next_pos, visited, doubled)) return part1, part2 if __name__ == "__main__": print(main()) print(f"{timeit(main, number=20) * 50:.3f} ms")