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))