aoc2021-12.py

import sys

EDGES = {tuple(line.strip().split("-")) for line in sys.stdin if line.strip()}
VERTICES = set(e[0] for e in EDGES) | set(e[1] for e in EDGES)
MAP = {
    v: set(
        w for w in VERTICES if (((v, w) in EDGES) or ((w, v) in EDGES)) and w != "start"
    )
    for v in VERTICES
}


def part1(v="start", visited=set()):
    if v == "end":
        return 1

    if v in visited and v.islower():
        return 0

    new_visited = visited | set([v])
    return sum(part1(w, new_visited) for w in MAP[v])


def part2(v="start", visited=set(), can_revisist_small=True):
    if v == "end":
        return 1

    is_small_cave_revisit = v in visited and v.islower()
    if is_small_cave_revisit and not can_revisist_small:
        return 0

    can_revisist_small = can_revisist_small and not is_small_cave_revisit
    new_visited = visited | set([v])
    return sum(part2(w, new_visited, can_revisist_small) for w in MAP[v])


print(f"Part 1: {part1()}")
print(f"Part 2: {part2()}")