97 lines
2.4 KiB
Python
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))
|