aocpy/2024/d16.py
2024-12-18 22:09:38 -05:00

97 lines
2.4 KiB
Python

import heapq
from collections import defaultdict
from lib import get_data
from lib import Grid2D
data = get_data(__file__)
g = Grid2D(data)
def search(start, end, dir):
paths = set()
def reconstruct(parents, current):
seen = set()
to_visit = [current]
while to_visit:
current = to_visit.pop()
if current in seen:
continue
seen.add(current)
paths.add(current[0])
for p in parents[current]:
to_visit.append(p)
def dist(n1, n2):
"""cost from node to node"""
if n1 == 0 or n1 == n2:
return 0
p1, d1 = n1
p2, d2 = n2
if p1 != p2:
return 1
if d1 != d2:
return 1000
assert False
def h(node):
"""heuristic function (never overestimate)"""
pos = node[0]
return abs(pos[0] - end[0]) + abs(pos[1] - end[1])
def is_goal(node):
return node[0] == end
def neighbors(node):
pos, dir = node
npos = pos[0] + dir[0], pos[1] + dir[1]
nbs = []
if g.contains(npos) and g[npos] != "#":
nbs.append((npos, dir))
ndir = dir[1], -dir[0]
nbs.append((pos, ndir))
ndir = -dir[1], dir[0]
nbs.append((pos, ndir))
return nbs
starts = [(start, dir)]
open_set = []
g_score = {}
cost = None
for start in starts:
heapq.heappush(open_set, (h(start), start))
g_score[start] = dist(0, start)
parents = defaultdict(list)
while open_set:
current_f_score, current = heapq.heappop(open_set)
if is_goal(current):
# assert current_f_score == g_score[current]
gs = g_score[current]
if cost is None or gs <= cost:
cost = gs
reconstruct(parents, current)
else:
break
for neighbor in neighbors(current):
tentative_g_score = g_score[current] + dist(current, neighbor)
if tentative_g_score <= g_score.get(neighbor, 10**12):
parents[neighbor].append(current)
g_score[neighbor] = tentative_g_score
f_score = g_score[neighbor] + h(neighbor)
heapq.heappush(open_set, (f_score, neighbor))
return cost, paths
(start,) = g.find("S")
(end,) = g.find("E")
dir = (0, 1)
cost, paths = search(start, end, dir)
print(cost)
print(len(paths))