~noelle/aoc-2021

19bdb9f828b8493680b1ef25e145a4d01e5457d6 — Noelle Leigh 2 years ago ab0e96c
12_2
2 files changed, 60 insertions(+), 0 deletions(-)

A 12/puzzle_2.py
A 12/small_example.txt
A 12/puzzle_2.py => 12/puzzle_2.py +53 -0
@@ 0,0 1,53 @@
"""
Solution for AoC 2021 12 Puzzle 2

cat input.txt | python puzzle_2.py
"""
import sys
from collections import Counter

START = "start"
END = "end"


def route(path: list[str], segments: list[set[str, str]]) -> int:
    """
    Return the number of unique paths to "end" starting with a given path.
    `path` will always have at least two caves in it.
    """
    head = path[-1]
    if head == END:
        return 1

    connecting_segments = filter(lambda segment: head in segment, segments)
    visited_small_caves = Counter(filter(lambda cave: cave.islower(), path))
    no_more_repeat_small_caves = visited_small_caves.most_common(1)[0][1] > 1
    path_total = 0
    for segment in connecting_segments:
        new_head = (segment - {head}).pop()
        # can't go back to the start
        if new_head == START:
            continue
        # Skip small caves that were already visited if a small cave has already
        # been visited twice
        if new_head.islower() and new_head in path and no_more_repeat_small_caves:
            continue
        path_total += route([*path, new_head], segments)

    return path_total


if __name__ == "__main__":
    cave_segments = list(
        map(
            lambda line: set(line.split("-", maxsplit=1)), sys.stdin.read().splitlines()
        )
    )
    distinct_path_count = 0
    # Loop through all caves that connect to start
    starting_segments = filter(lambda segment: START in segment, cave_segments)
    for starting_segment in starting_segments:
        second_cave = (starting_segment - {START}).pop()
        distinct_path_count += route([START, second_cave], cave_segments)

    sys.stdout.write(str(distinct_path_count))

A 12/small_example.txt => 12/small_example.txt +7 -0
@@ 0,0 1,7 @@
start-A
start-b
A-c
A-b
b-d
A-end
b-end