75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
|
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()
|
||
|
|
||
|
|
||
|
|
||
|
|
||
|
|