astar.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. """Shortest paths and path lengths using the A* ("A star") algorithm.
  2. """
  3. from heapq import heappop, heappush
  4. from itertools import count
  5. import networkx as nx
  6. from networkx.algorithms.shortest_paths.weighted import _weight_function
  7. __all__ = ["astar_path", "astar_path_length"]
  8. def astar_path(G, source, target, heuristic=None, weight="weight"):
  9. """Returns a list of nodes in a shortest path between source and target
  10. using the A* ("A-star") algorithm.
  11. There may be more than one shortest path. This returns only one.
  12. Parameters
  13. ----------
  14. G : NetworkX graph
  15. source : node
  16. Starting node for path
  17. target : node
  18. Ending node for path
  19. heuristic : function
  20. A function to evaluate the estimate of the distance
  21. from the a node to the target. The function takes
  22. two nodes arguments and must return a number.
  23. If the heuristic is inadmissible (if it might
  24. overestimate the cost of reaching the goal from a node),
  25. the result may not be a shortest path.
  26. The algorithm does not support updating heuristic
  27. values for the same node due to caching the first
  28. heuristic calculation per node.
  29. weight : string or function
  30. If this is a string, then edge weights will be accessed via the
  31. edge attribute with this key (that is, the weight of the edge
  32. joining `u` to `v` will be ``G.edges[u, v][weight]``). If no
  33. such edge attribute exists, the weight of the edge is assumed to
  34. be one.
  35. If this is a function, the weight of an edge is the value
  36. returned by the function. The function must accept exactly three
  37. positional arguments: the two endpoints of an edge and the
  38. dictionary of edge attributes for that edge. The function must
  39. return a number or None to indicate a hidden edge.
  40. Raises
  41. ------
  42. NetworkXNoPath
  43. If no path exists between source and target.
  44. Examples
  45. --------
  46. >>> G = nx.path_graph(5)
  47. >>> print(nx.astar_path(G, 0, 4))
  48. [0, 1, 2, 3, 4]
  49. >>> G = nx.grid_graph(dim=[3, 3]) # nodes are two-tuples (x,y)
  50. >>> nx.set_edge_attributes(G, {e: e[1][0] * 2 for e in G.edges()}, "cost")
  51. >>> def dist(a, b):
  52. ... (x1, y1) = a
  53. ... (x2, y2) = b
  54. ... return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5
  55. >>> print(nx.astar_path(G, (0, 0), (2, 2), heuristic=dist, weight="cost"))
  56. [(0, 0), (0, 1), (0, 2), (1, 2), (2, 2)]
  57. Notes
  58. -----
  59. Edge weight attributes must be numerical.
  60. Distances are calculated as sums of weighted edges traversed.
  61. The weight function can be used to hide edges by returning None.
  62. So ``weight = lambda u, v, d: 1 if d['color']=="red" else None``
  63. will find the shortest red path.
  64. See Also
  65. --------
  66. shortest_path, dijkstra_path
  67. """
  68. if source not in G or target not in G:
  69. msg = f"Either source {source} or target {target} is not in G"
  70. raise nx.NodeNotFound(msg)
  71. if heuristic is None:
  72. # The default heuristic is h=0 - same as Dijkstra's algorithm
  73. def heuristic(u, v):
  74. return 0
  75. push = heappush
  76. pop = heappop
  77. weight = _weight_function(G, weight)
  78. # The queue stores priority, node, cost to reach, and parent.
  79. # Uses Python heapq to keep in priority order.
  80. # Add a counter to the queue to prevent the underlying heap from
  81. # attempting to compare the nodes themselves. The hash breaks ties in the
  82. # priority and is guaranteed unique for all nodes in the graph.
  83. c = count()
  84. queue = [(0, next(c), source, 0, None)]
  85. # Maps enqueued nodes to distance of discovered paths and the
  86. # computed heuristics to target. We avoid computing the heuristics
  87. # more than once and inserting the node into the queue too many times.
  88. enqueued = {}
  89. # Maps explored nodes to parent closest to the source.
  90. explored = {}
  91. while queue:
  92. # Pop the smallest item from queue.
  93. _, __, curnode, dist, parent = pop(queue)
  94. if curnode == target:
  95. path = [curnode]
  96. node = parent
  97. while node is not None:
  98. path.append(node)
  99. node = explored[node]
  100. path.reverse()
  101. return path
  102. if curnode in explored:
  103. # Do not override the parent of starting node
  104. if explored[curnode] is None:
  105. continue
  106. # Skip bad paths that were enqueued before finding a better one
  107. qcost, h = enqueued[curnode]
  108. if qcost < dist:
  109. continue
  110. explored[curnode] = parent
  111. for neighbor, w in G[curnode].items():
  112. cost = weight(curnode, neighbor, w)
  113. if cost is None:
  114. continue
  115. ncost = dist + cost
  116. if neighbor in enqueued:
  117. qcost, h = enqueued[neighbor]
  118. # if qcost <= ncost, a less costly path from the
  119. # neighbor to the source was already determined.
  120. # Therefore, we won't attempt to push this neighbor
  121. # to the queue
  122. if qcost <= ncost:
  123. continue
  124. else:
  125. h = heuristic(neighbor, target)
  126. enqueued[neighbor] = ncost, h
  127. push(queue, (ncost + h, next(c), neighbor, ncost, curnode))
  128. raise nx.NetworkXNoPath(f"Node {target} not reachable from {source}")
  129. def astar_path_length(G, source, target, heuristic=None, weight="weight"):
  130. """Returns the length of the shortest path between source and target using
  131. the A* ("A-star") algorithm.
  132. Parameters
  133. ----------
  134. G : NetworkX graph
  135. source : node
  136. Starting node for path
  137. target : node
  138. Ending node for path
  139. heuristic : function
  140. A function to evaluate the estimate of the distance
  141. from the a node to the target. The function takes
  142. two nodes arguments and must return a number.
  143. If the heuristic is inadmissible (if it might
  144. overestimate the cost of reaching the goal from a node),
  145. the result may not be a shortest path.
  146. The algorithm does not support updating heuristic
  147. values for the same node due to caching the first
  148. heuristic calculation per node.
  149. weight : string or function
  150. If this is a string, then edge weights will be accessed via the
  151. edge attribute with this key (that is, the weight of the edge
  152. joining `u` to `v` will be ``G.edges[u, v][weight]``). If no
  153. such edge attribute exists, the weight of the edge is assumed to
  154. be one.
  155. If this is a function, the weight of an edge is the value
  156. returned by the function. The function must accept exactly three
  157. positional arguments: the two endpoints of an edge and the
  158. dictionary of edge attributes for that edge. The function must
  159. return a number or None to indicate a hidden edge.
  160. Raises
  161. ------
  162. NetworkXNoPath
  163. If no path exists between source and target.
  164. See Also
  165. --------
  166. astar_path
  167. """
  168. if source not in G or target not in G:
  169. msg = f"Either source {source} or target {target} is not in G"
  170. raise nx.NodeNotFound(msg)
  171. weight = _weight_function(G, weight)
  172. path = astar_path(G, source, target, heuristic, weight)
  173. return sum(weight(u, v, G[u][v]) for u, v in zip(path[:-1], path[1:]))