[Python3] Dijkstra's algorithm with 14 lines
Yesterday I implemented Prime factorization in 14 lines, but today I implemented Dijkstra's algorithm in 14 lines using the heap. In 14 lines from the top.
dijkstra.py
from heapq import heappush, heappop
def dijkstra(vertices_num, edges, src):
dist = [float('inf')] * vertices_num
heap = [(0, src)]
while len(heap):
d, from_ = heappop(heap)
if d < dist[from_]:
dist[from_] = d
for _, to, cost in filter(lambda e: e[0]==from_, edges):
heappush(heap, (min(dist[to], dist[from_] + cost), to))
return dist
if __name__ == '__main__':
answer = dijkstra(
8,
[
(0, 1, 5), (0, 2, 4), (0, 3, 1), (1, 4, 2), (1, 0, 5), (2, 0, 4), (2, 3, 2), (2, 4, 5),
(2, 5, 6), (3, 0, 1), (3, 2, 2), (4, 1, 2), (4, 2, 5), (4, 6, 1), (4, 7, 3), (5, 2, 6),
(5, 7, 2), (6, 4, 1), (6, 7, 4), (7, 4, 3), (7, 5, 2), (7, 6, 4)
], 0
)
print(answer) # [0, 5, 3, 1, 7, 9, 8, 10]
The first argument is the number of vertices, the second argument is a list of (from_, to, cost), and the third argument is the index of the starting point. The return value returns the shortest distance from the starting point to each vertex.
reference
https://www.youtube.com/watch?v=X1AsMlJdiok
Regarding the Dijkstra method, the explanation in this video was very easy to understand, so I made a test case with reference to the problems dealt with there. Red letters are the indexes of each vertex.
(Source: From the Youtube page above)
Recommended Posts