85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
"""A* shortest-path search on a weighted graph."""
|
|
from __future__ import annotations
|
|
import heapq
|
|
from typing import Callable, Dict, Hashable, List, Tuple
|
|
|
|
Node = Hashable
|
|
Graph = Dict[Node, Dict[Node, float]]
|
|
|
|
|
|
def astar(
|
|
graph: Graph,
|
|
start: Node,
|
|
goal: Node,
|
|
h: Callable[[Node], float] = None,
|
|
) -> Tuple[float, List[Node]]:
|
|
if start not in graph:
|
|
raise ValueError(f"start node {start!r} not in graph")
|
|
if goal not in graph:
|
|
raise ValueError(f"goal node {goal!r} not in graph")
|
|
if h is None:
|
|
h = lambda _n: 0.0
|
|
|
|
g: Dict[Node, float] = {start: 0.0}
|
|
came: Dict[Node, Node] = {}
|
|
closed: set = set()
|
|
counter = 0
|
|
open_heap: list = [(h(start), counter, start)]
|
|
counter += 1
|
|
|
|
while open_heap:
|
|
_f, _c, current = heapq.heappop(open_heap)
|
|
if current in closed:
|
|
continue
|
|
if current == goal:
|
|
path = [current]
|
|
while current in came:
|
|
current = came[current]
|
|
path.append(current)
|
|
path.reverse()
|
|
return g[goal], path
|
|
closed.add(current)
|
|
for nbr, w in graph.get(current, {}).items():
|
|
if nbr in closed:
|
|
continue
|
|
if w < 0:
|
|
raise ValueError("negative edge weights are not supported")
|
|
tentative = g[current] + w
|
|
if tentative < g.get(nbr, float("inf")):
|
|
g[nbr] = tentative
|
|
came[nbr] = current
|
|
f = tentative + h(nbr)
|
|
heapq.heappush(open_heap, (f, counter, nbr))
|
|
counter += 1
|
|
|
|
return float("inf"), []
|
|
|
|
|
|
def dijkstra(graph, start, goal):
|
|
if start not in graph or goal not in graph:
|
|
raise ValueError("node not in graph")
|
|
dist = {start: 0.0}
|
|
prev = {}
|
|
pq = [(0.0, start)]
|
|
visited = set()
|
|
while pq:
|
|
d, u = heapq.heappop(pq)
|
|
if u in visited:
|
|
continue
|
|
visited.add(u)
|
|
if u == goal:
|
|
break
|
|
for v, w in graph.get(u, {}).items():
|
|
nd = d + w
|
|
if nd < dist.get(v, float("inf")):
|
|
dist[v] = nd
|
|
prev[v] = u
|
|
heapq.heappush(pq, (nd, v))
|
|
if goal not in dist:
|
|
return float("inf"), []
|
|
path = [goal]
|
|
while path[-1] != start:
|
|
path.append(prev[path[-1]])
|
|
path.reverse()
|
|
return dist[goal], path
|