GraphAlgo.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171
  1. """
  2. altgraph.GraphAlgo - Graph algorithms
  3. =====================================
  4. """
  5. from altgraph import GraphError
  6. def dijkstra(graph, start, end=None):
  7. """
  8. Dijkstra's algorithm for shortest paths
  9. `David Eppstein, UC Irvine, 4 April 2002
  10. <http://www.ics.uci.edu/~eppstein/161/python/>`_
  11. `Python Cookbook Recipe
  12. <http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/119466>`_
  13. Find shortest paths from the start node to all nodes nearer than or
  14. equal to the end node.
  15. Dijkstra's algorithm is only guaranteed to work correctly when all edge
  16. lengths are positive. This code does not verify this property for all
  17. edges (only the edges examined until the end vertex is reached), but will
  18. correctly compute shortest paths even for some graphs with negative edges,
  19. and will raise an exception if it discovers that a negative edge has
  20. caused it to make a mistake.
  21. Adapted to altgraph by Istvan Albert, Pennsylvania State University -
  22. June, 9 2004
  23. """
  24. D = {} # dictionary of final distances
  25. P = {} # dictionary of predecessors
  26. Q = _priorityDictionary() # estimated distances of non-final vertices
  27. Q[start] = 0
  28. for v in Q:
  29. D[v] = Q[v]
  30. if v == end:
  31. break
  32. for w in graph.out_nbrs(v):
  33. edge_id = graph.edge_by_node(v, w)
  34. vwLength = D[v] + graph.edge_data(edge_id)
  35. if w in D:
  36. if vwLength < D[w]:
  37. raise GraphError(
  38. "Dijkstra: found better path to already-final vertex"
  39. )
  40. elif w not in Q or vwLength < Q[w]:
  41. Q[w] = vwLength
  42. P[w] = v
  43. return (D, P)
  44. def shortest_path(graph, start, end):
  45. """
  46. Find a single shortest path from the *start* node to the *end* node.
  47. The input has the same conventions as dijkstra(). The output is a list of
  48. the nodes in order along the shortest path.
  49. **Note that the distances must be stored in the edge data as numeric data**
  50. """
  51. D, P = dijkstra(graph, start, end)
  52. Path = []
  53. while 1:
  54. Path.append(end)
  55. if end == start:
  56. break
  57. end = P[end]
  58. Path.reverse()
  59. return Path
  60. #
  61. # Utility classes and functions
  62. #
  63. class _priorityDictionary(dict):
  64. """
  65. Priority dictionary using binary heaps (internal use only)
  66. David Eppstein, UC Irvine, 8 Mar 2002
  67. Implements a data structure that acts almost like a dictionary, with
  68. two modifications:
  69. 1. D.smallest() returns the value x minimizing D[x]. For this to
  70. work correctly, all values D[x] stored in the dictionary must be
  71. comparable.
  72. 2. iterating "for x in D" finds and removes the items from D in sorted
  73. order. Each item is not removed until the next item is requested,
  74. so D[x] will still return a useful value until the next iteration
  75. of the for-loop. Each operation takes logarithmic amortized time.
  76. """
  77. def __init__(self):
  78. """
  79. Initialize priorityDictionary by creating binary heap of pairs
  80. (value,key). Note that changing or removing a dict entry will not
  81. remove the old pair from the heap until it is found by smallest()
  82. or until the heap is rebuilt.
  83. """
  84. self.__heap = []
  85. dict.__init__(self)
  86. def smallest(self):
  87. """
  88. Find smallest item after removing deleted items from front of heap.
  89. """
  90. if len(self) == 0:
  91. raise IndexError("smallest of empty priorityDictionary")
  92. heap = self.__heap
  93. while heap[0][1] not in self or self[heap[0][1]] != heap[0][0]:
  94. lastItem = heap.pop()
  95. insertionPoint = 0
  96. while 1:
  97. smallChild = 2 * insertionPoint + 1
  98. if (
  99. smallChild + 1 < len(heap)
  100. and heap[smallChild] > heap[smallChild + 1]
  101. ):
  102. smallChild += 1
  103. if smallChild >= len(heap) or lastItem <= heap[smallChild]:
  104. heap[insertionPoint] = lastItem
  105. break
  106. heap[insertionPoint] = heap[smallChild]
  107. insertionPoint = smallChild
  108. return heap[0][1]
  109. def __iter__(self):
  110. """
  111. Create destructive sorted iterator of priorityDictionary.
  112. """
  113. def iterfn():
  114. while len(self) > 0:
  115. x = self.smallest()
  116. yield x
  117. del self[x]
  118. return iterfn()
  119. def __setitem__(self, key, val):
  120. """
  121. Change value stored in dictionary and add corresponding pair to heap.
  122. Rebuilds the heap if the number of deleted items gets large, to avoid
  123. memory leakage.
  124. """
  125. dict.__setitem__(self, key, val)
  126. heap = self.__heap
  127. if len(heap) > 2 * len(self):
  128. self.__heap = [(v, k) for k, v in self.items()]
  129. self.__heap.sort()
  130. else:
  131. newPair = (val, key)
  132. insertionPoint = len(heap)
  133. heap.append(None)
  134. while insertionPoint > 0 and newPair < heap[(insertionPoint - 1) // 2]:
  135. heap[insertionPoint] = heap[(insertionPoint - 1) // 2]
  136. insertionPoint = (insertionPoint - 1) // 2
  137. heap[insertionPoint] = newPair
  138. def setdefault(self, key, val):
  139. """
  140. Reimplement setdefault to pass through our customized __setitem__.
  141. """
  142. if key not in self:
  143. self[key] = val
  144. return self[key]