aocpy/2021/d15.py
2024-11-23 14:00:08 -05:00

84 lines
2.2 KiB
Python

from lib import get_data
from lib import Grid2D
import heapq
data = get_data(__file__)
def new_data(data, repeat=5):
new_lines = []
for line in data.splitlines():
new_line = str(line)
for i in range(len(line) * repeat):
ni = i - len(line)
if ni < 0:
continue
nc = int(new_line[ni]) + 1
if nc == 10:
nc = 1
new_line += str(nc)
new_lines.append(new_line)
for i in range(len(data.splitlines()) * repeat):
ni = i - len(data.splitlines())
if ni < 0:
continue
line = list(new_lines[ni])
for i in range(len(line)):
nc = int(line[i]) + 1
if nc == 10:
nc = 1
line[i] = str(nc)
line = "".join(line)
new_lines.append(line)
data = "\n".join(new_lines)
return data
big_data = new_data(data)
for data in [data, big_data]:
g = Grid2D(data)
def dist(n1, n2):
"""cost from node to node"""
if n1 == 0:
return 0
return int(g[n2])
def h(node):
"""heuristic function (never overestimate)"""
return abs(g.n_rows - 1 - node[0]) + abs(g.n_cols - 1 - node[1])
def is_goal(node):
return node == (g.n_rows - 1, g.n_cols - 1)
def neighbors(node):
return list(g.neighbors_ort(node))
starts = [(0, 0)]
open_set = []
g_score = {}
cost = None
for start in starts:
heapq.heappush(open_set, (h(start), start))
g_score[start] = dist(0, start)
while open_set:
current_f_score, current = heapq.heappop(open_set)
if is_goal(current):
assert current_f_score == g_score[current]
cost = g_score[current]
break
for neighbor in neighbors(current):
tentative_g_score = g_score[current] + dist(current, neighbor)
if neighbor not in g_score or tentative_g_score < g_score[neighbor]:
g_score[neighbor] = tentative_g_score
f_score = g_score[neighbor] + h(neighbor)
heapq.heappush(open_set, (f_score, neighbor))
print(cost)