Files
python-aoc-2021/day15/part2.py
Sebastian Seedorf 18c76f694b Day 15 (speedup)
2021-12-16 10:03:16 +01:00

32 lines
1.1 KiB
Python

#!/usr/bin/env python3
import heapq
import numpy as np
arr = np.array([[int(n) for n in line.strip()] for line in open("input.txt")], dtype='int16')
add = np.array([list(range(i, i+5)) for i in range(5)], dtype='int16')
add = np.repeat(np.repeat(add, arr.shape[0], axis=0), arr.shape[1], axis=1)
arr = (np.tile(arr, (5, 5)) + add - 1) % 9 + 1
open_list = [(arr[0, 0], (0, 0))]
closed_list = set()
target = ty, tx = arr[:-1, :-1].shape
while len(open_list):
val, curr = heapq.heappop(open_list)
if curr in closed_list:
continue
if curr == target:
print(val - arr[0, 0])
break
y, x = curr
closed_list.add(curr)
if 0 < y and (nxt := (y-1, x)) not in closed_list:
heapq.heappush(open_list, (arr[nxt] + val, nxt))
if 0 < x and (nxt := (y, x-1)) not in closed_list:
heapq.heappush(open_list, (arr[nxt] + val, nxt))
if ty > y and (nxt := (y+1, x)) not in closed_list:
heapq.heappush(open_list, (arr[nxt] + val, nxt))
if tx > x and (nxt := (y, x+1)) not in closed_list:
heapq.heappush(open_list, (arr[nxt] + val, nxt))