Advent of Code 2024 Day 16

from pathlib import Path
from collections import deque
from pprint import pprint

TEST = """**REDACTED**"""

DIRECTIONS = (1, -1, 1j, -1j)


def parse(input: str) -> tuple[complex, complex, set[complex]]:
    walls: set[complex] = set()
    start: complex = 0j
    end: complex = 0j
    for y, row in enumerate(input.splitlines()):
        for x, c in enumerate(row):
            match c:
                case "#":
                    walls.add(complex(x, y))
                case "S":
                    start = complex(x, y)
                case "E":
                    end = complex(x, y)
                case _:
                    pass
    return start, end, walls


def walk(
    start: complex, end: complex, walls: set[complex]
) -> dict[complex, dict[complex, int]]:
    scores: dict[complex, dict[complex, int]] = {1: {start: 0}, -1: {}, 1j: {}, -1j: {}}
    points: deque[tuple[complex, complex]] = deque([(start, 1)])
    while points:
        p, d = points.popleft()
        score = scores[d][p]
        for move, dir, cost in [(p + d, d, 1), (p, d * -1j, 1000), (p, d * 1j, 1000)]:
            if move in walls:
                continue
            if move in scores[dir] and scores[dir][move] <= score + cost:
                continue
            scores[dir][move] = score + cost
            points.append((move, dir))
    return scores


def part1(input: str) -> int:
    start, end, walls = parse(input)
    scores = walk(start, end, walls)
    return min(scores[d][end] for d in DIRECTIONS if end in scores[d])


assert part1(TEST) == 7036

INPUT = Path("input/day16.txt").read_text()
part1_total = part1(INPUT)
print(f"{part1_total=:,}")  # 106,512


def part2(input: str) -> int:
    start, end, walls = parse(input)
    scores = walk(start, end, walls)
    path: set[complex] = {start, end}
    cost = min(scores[d][end] for d in DIRECTIONS if end in scores[d])
    points: deque[complex, complex, int] = deque(
        [
            (end, d, cost)
            for d in DIRECTIONS
            if end in scores[d] and scores[d][end] == cost
        ]
    )
    while points:
        p, d, cost = points.popleft()
        for move, dir, c1 in [
            (p - d, d, cost - 1),
            (p, d * 1j, cost - 1000),
            (p, d * -1j, cost - 1000),
        ]:
            if move in scores[dir] and scores[dir][move] == c1:
                path.add(move)
                points.append((move, dir, c1))
    return len(path)


assert part2(TEST) == 45
part2_total = part2(INPUT)
print(f"{part2_total=:,}")  # 563