79 lines
2.3 KiB
Python
79 lines
2.3 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]
|
|
|
|
(w, h) = (len(map[0]), len(map))
|
|
|
|
heap: list[tuple[int, tuple[int, int], int, int, tuple[Any, ...]]] = [
|
|
# cost, pos, dir, straight, trace
|
|
(0, (0, 0), 0, 0, ()),
|
|
(0, (0, 0), 3, 0, ()),
|
|
]
|
|
seen = set()
|
|
it = 0
|
|
while heap:
|
|
it += 1
|
|
(cost, (x, y), dir, straight, trace) = heapq.heappop(heap)
|
|
if it >= 10000:
|
|
if (x, y, dir, straight) in seen:
|
|
continue
|
|
seen.add((x, y, dir, straight))
|
|
|
|
if (x, y) == (len(map[0]) - 1, len(map) - 1) and 10 >= straight > 3:
|
|
break
|
|
|
|
(dx, dy) = DIRS[dir]
|
|
if (
|
|
10 >= (straight + 1)
|
|
and (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)
|
|
)
|
|
|
|
if 11 >= (straight + 1) > 4:
|
|
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, 1, 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, 1, newtrace))
|
|
else:
|
|
print("oh no")
|
|
sys.exit(1)
|
|
|
|
print()
|
|
|
|
path = [[str(c) for c in row] for row in map]
|
|
for (x, y, dir) in trace:
|
|
# path[y][x] = "\033[93m" + [">", "^", "<", "v"][dir] + "\033[0m"
|
|
path[y][x] = f"\033[93m{path[y][x]}\033[0m"
|
|
print("\n".join("".join(line) for line in path))
|
|
print(cost)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
# 1040 too low
|
|
# 1066 too high
|