AOC2021 Day 12 (Python)

day12_tests.py

def cave_system(lines):
    edges = {}
    for line in lines:
        start_label, end_label = line.strip().split("-")

        targets = edges.get(start_label, set())
        targets.add(end_label)
        edges[start_label] = targets

        reverse_targets = edges.get(end_label, set())
        reverse_targets.add(start_label)
        edges[end_label] = reverse_targets

    return edges


def paths_to_end(edges, path, has_double_visited_small, visited_small, current_label):
    if current_label == 'end':
        yield *path, 'end'
        return

    possible_targets = edges.get(current_label, set())
    possible_targets.discard('start')
    targets = possible_targets - visited_small if has_double_visited_small else possible_targets

    is_small = all(str.islower(c) for c in current_label)

    if is_small:
        new_visited = set(visited_small)
        new_visited.add(current_label)
    else:
        new_visited = visited_small

    for target in targets:
        new_has_double_visited_small = has_double_visited_small or target in new_visited
        for found_path in paths_to_end(edges, (*path, current_label), new_has_double_visited_small, new_visited, target):
            yield found_path


def test_count_paths():
    example = [
        'start-A',
        'start-b',
        'A-c',
        'A-b',
        'b-d',
        'A-end',
        'b-end'
    ]

    example_edges = cave_system(example)
    assert sum(1 for path in paths_to_end(example_edges, (), False, set(), 'start')) == 36

    with open("puzzle_inputs/day12.txt") as file:
        edges = cave_system(file)

    print()
    print(sum(1 for path in paths_to_end(edges, (), True, set(), 'start')))
    print(sum(1 for path in paths_to_end(edges, (), False, set(), 'start')))