2023/day17.py

import heapq
from pprint import pprint
import sys
from typing import Any


DIRS = [(1, 0), (0, -1), (-1, 0), (0, 1)]


def main() -> None:
    with open(sys.argv[1]) as f:
        map = [[int(c) for c in l.strip()] for l in f]

    pprint(map)

    heap: list[tuple[int, tuple[int, int], int, int, tuple[Any, ...]]] = [
        (0, (0, 0), 0, 0, ())
    ]
    seen = set()
    while heap:
        (cost, (x, y), dir, straight, trace) = heapq.heappop(heap)
        if (x, y, dir, straight) in seen:
            continue
        seen.add((x, y, dir, straight))

        if straight >= 3:
            continue
        if (x, y) == (len(map) - 1, len(map[0]) - 1):
            break

        (dx, dy) = DIRS[dir]
        if (0 <= (y + dy) < len(map)) and (0 <= (x + dx) < len(map[y])):
            newtrace = (*trace, (x + dx, y + dy, dir))
            newcost = cost + map[y + dy][x + dx]
            heapq.heappush(
                heap, (newcost, (x + dx, y + dy), dir, straight + 1, newtrace)
            )

        newdir = (dir + 1) % 4
        (dx, dy) = DIRS[newdir]
        if (0 <= (y + dy) < len(map)) and (0 <= (x + dx) < len(map[y])):
            newtrace = (*trace, (x + dx, y + dy, newdir))
            newcost = cost + map[y + dy][x + dx]
            heapq.heappush(heap, (newcost, (x + dx, y + dy), newdir, 0, newtrace))

        newdir = (dir + 3) % 4
        (dx, dy) = DIRS[newdir]
        if (0 <= (y + dy) < len(map)) and (0 <= (x + dx) < len(map[y])):
            newtrace = (*trace, (x + dx, y + dy, newdir))
            newcost = cost + map[y + dy][x + dx]
            heapq.heappush(heap, (newcost, (x + dx, y + dy), newdir, 0, newtrace))
    else:
        print("oh no")
        path = [["."] * len(map[y]) for y in range(len(map))]
        for (x, y, dir) in trace:
            path[y][x] = [">", "^", "<", "v"][dir]
        print("\n".join("".join(line) for line in path))
        sys.exit(1)

    print()

    path = [["."] * len(map[y]) for y in range(len(map))]
    for (x, y, dir) in trace:
        path[y][x] = [">", "^", "<", "v"][dir]
    print("\n".join("".join(line) for line in path))

    cost = 0
    for (x, y, _) in trace:
        cost += map[y][x]
    print(cost)


if __name__ == "__main__":
    main()