#!/usr/bin/env python
import numpy as np
import math
from heapq import *
from collections import defaultdict
with open('day15.txt') as data:
costs = []
for line in data:
costs.append(list(line.strip()))
costs = np.array(costs).astype(int)
def a_star(costs) -> int:
def h(point):
m, n = point
# approximate with straight line distance between points
return math.sqrt((costs.shape[0]-m)**2 + (costs.shape[1]-n)**2)
def neighbors(point):
m, n = point
neighboring = {(0, 1), (1, 0), (0, -1), (-1, 0)}
return {(m+dx, n+dy) for dx, dy in neighboring
if 0<=m+dx<costs.shape[0] and 0<=n+dy<costs.shape[1]}
def path_sum(parents, current):
complete = [current]
while current in parents:
current = parents[current]
if current != (0, 0):
complete.insert(0, current)
return sum(map(lambda x: costs[x], complete))
discovered = [(0,0)]
parents = {}
gScore = defaultdict(lambda : math.inf)
gScore[(0,0)] = 0
fScore = defaultdict(lambda : math.inf)
fScore[(0,0)] = h((0,0))
while discovered:
if discovered[0] == (goal := (costs.shape[0]-1, costs.shape[1]-1)):
return path_sum(parents, goal)
current = heappop(discovered)
for neighbor in neighbors(current):
tentative = gScore[current] + costs[neighbor]
if tentative < gScore[neighbor]:
parents[neighbor] = current
gScore[neighbor] = tentative
fScore[neighbor] = tentative + h(neighbor)
if neighbor not in discovered:
heappush(discovered, neighbor)
else:
return None
# part 1
print(a_star(costs))
# part 2
new_costs = np.zeros((costs.shape[0] * 5, costs.shape[1] * 5)).astype(int)
def increment(arr):
arr += 1
arr[arr > 9] = 1
xDim, yDim = costs.shape
for i in range(5):
for j in range(5):
copied = np.copy(costs)
for _ in range(i+j):
increment(copied)
new_costs[xDim*i:xDim*(i+1), yDim*j:yDim*(j+1)] = copied
print(a_star(new_costs))