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()