aoc2024day18.py

# === Parsing ===
from itertools import product

SIZE = 70

with open("input.txt", "r") as f:
    falling_bytes = [tuple(int(num) for num in line.split(",")) for line in f]


valid = set(product(range(SIZE + 1), repeat=2))
start = (0, 0)
end = (SIZE, SIZE)


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


DIRECTIONS = [
    (1, 0),
    (0, 1),
    (-1, 0),
    (0, -1),
]
FIRST_BYTES = 1024


def find_path(start, end, valid):
    """A* search algorithm."""
    ex, ey = end
    costs = defaultdict(lambda: float("inf"))
    costs[start] = 0
    heap = [(2 * SIZE, start)]  # Pairs of (heuristic cost, pos)
    while heap:
        _, pos = heapq.heappop(heap)
        if pos == end:
            return costs[pos]

        px, py = pos
        cost = costs[pos]
        for dx, dy in DIRECTIONS:
            nx, ny = px + dx, py + dy
            if (nx, ny) in valid and cost + 1 < costs[nx, ny]:
                costs[nx, ny] = cost + 1
                hcost = costs[nx, ny] + abs(ex - nx) + abs(ey - ny)
                heapq.heappush(heap, (hcost, (nx, ny)))

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


valid_p1 = valid - set(falling_bytes[:FIRST_BYTES])
cost = find_path(start, end, valid_p1)
print(cost)


# === Part 2 ===

valid_p2 = set(valid_p1)
for falling_byte in falling_bytes[FIRST_BYTES:]:
    valid_p2.remove(falling_byte)
    try:
        find_path(start, end, valid_p2)
    except RuntimeError:
        print(falling_byte)
        break