aoc2024day16.py

# === Parsing ===

start = None
end = None
valid = set()
with open("input.txt", "r") as f:
    for y, line in enumerate(f):
        for x, char in enumerate(line.strip()):
            if char in ".SE":
                valid.add((x, y))
                if char == "S":
                    start = (x, y)
                elif char == "E":
                    end = (x, y)

assert start is not None
assert end is not None

print(start)


# === Part 1 ===
from collections import Counter, defaultdict
import heapq

# Dijkstra's algorithm


def find_path_cost(start, end, valid):
    # Heap of the cells we're exploring with their cheapest costs
    heap = []
    heapq.heapify(heap)
    # Maps positions to the cheapest path so far.
    costs = defaultdict(lambda: float("inf"))
    costs[start, (1, 0)] = 0
    heapq.heappush(heap, (0, start, (1, 0)))  # Triples of (cost, position, (dx, dy))
    while heap:
        cost, (px, py), (dx, dy) = heapq.heappop(heap)
        if (px, py) == end:
            return cost

        next_possible_moves = [
            (cost + 1, (px + dx, py + dy), (dx, dy)),  # Move forward
            (cost + 1000, (px, py), (-dy, -dx)),  # Turn left
            (cost + 1000, (px, py), (dy, dx)),  # Turn right
        ]
        for ncost, (nx, ny), (ndx, ndy) in next_possible_moves:
            if (nx, ny) in valid and ncost < costs[(nx, ny), (ndx, ndy)]:
                costs[(nx, ny), (ndx, ndy)] = ncost
                heapq.heappush(heap, (ncost, (nx, ny), (ndx, ndy)))

    raise RuntimeError(f"Couldn't reach {end} from {start}.")


print(find_path_cost(start, end, valid))


# === Part 1, alternative ===
def alt_find_path_cost(start, end, valid):
    """Slower but simpler."""
    heap = []
    heapq.heappush(heap, (0, start, (1, 0)))
    visited = set()
    while heap:
        cost, (px, py), (dx, dy) = heapq.heappop(heap)
        visited.add(((px, py), (dx, dy)))
        if (px, py) == end:
            return cost

        next_possible_moves = [
            (cost + 1, (px + dx, py + dy), (dx, dy)),  # Move forward
            (cost + 1000, (px, py), (-dy, -dx)),  # Turn left
            (cost + 1000, (px, py), (dy, dx)),  # Turn right
        ]
        for ncost, (nx, ny), (ndx, ndy) in next_possible_moves:
            if (nx, ny) in valid and ((nx, ny), (ndx, ndy)) not in visited:
                heapq.heappush(heap, (ncost, (nx, ny), (ndx, ndy)))

    raise RuntimeError(f"Couldn't reach {end} from {start}.")


print(alt_find_path_cost(start, end, valid))


# === Part 2 ===
def count_cells_in_better_paths(start, end, valid, best_cost):
    # Heap of the cells we're exploring with their cheapest costs
    heap = []
    heapq.heapify(heap)
    # Maps positions to the cheapest path so far.
    costs = defaultdict(lambda: float("inf"))
    costs[start, (1, 0)] = 0
    heapq.heappush(
        heap, (0, start, (1, 0), {start})
    )  # Quadruples of (cost, position, (dx, dy), cells_in_the_path)
    all_visited = set()  # Set of all cells visited on a short path to the end.
    while heap:
        cost, (px, py), (dx, dy), seen = heapq.heappop(heap)
        if (px, py) == end:
            all_visited |= seen
            continue

        nx, ny = px + dx, py + dy
        next_possible_moves = [
            (cost + 1, (nx, ny), (dx, dy), seen | {(nx, ny)}),  # Move forward
            (cost + 1000, (px, py), (-dy, -dx), seen),  # Turn left
            (cost + 1000, (px, py), (dy, dx), seen),  # Turn right
        ]
        for ncost, (nx, ny), (ndx, ndy), nseen in next_possible_moves:
            if (
                (nx, ny) in valid
                and ncost <= costs[(nx, ny), (ndx, ndy)]
                and ncost <= best_cost
            ):
                costs[(nx, ny), (ndx, ndy)] = ncost
                heapq.heappush(heap, (ncost, (nx, ny), (ndx, ndy), nseen))

    return len(all_visited)


best_cost = find_path_cost(start, end, valid)
print(count_cells_in_better_paths(start, end, valid, best_cost))