[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. image.png (Source: From the Youtube page above)

Recommended Posts

[Python3] Dijkstra's algorithm with 14 lines
Implementation of Dijkstra's algorithm with python
Algorithm in Python (Dijkstra's algorithm)
Find the shortest path with the Python Dijkstra's algorithm
Implement Dijkstra's Algorithm in python
Python algorithm
FizzBuzz with Python3
Algorithm learned with Python 10th: Binary search
Scraping with Python
Algorithm learned with Python 5th: Fibonacci sequence
Algorithm learned with Python 9th: Linear search
Statistics with python
Python memorandum (algorithm)
Scraping with Python
Ant book in python: Sec.2-5 Dijkstra's algorithm
Python with Go
Algorithm learned with Python 7th: Year conversion
Algorithm learned with Python 8th: Evaluation of algorithm
Twilio with Python
Algorithm learned with Python 4th: Prime numbers
Search the maze with the python A * algorithm
Play with 2016-Python
Tested with Python
Algorithm learned with Python 2nd: Vending machine
Algorithm learned with Python 19th: Sorting (heapsort)
with syntax (Python)
Algorithm learned with Python 6th: Leap year
Bingo with python
Zundokokiyoshi with python
Algorithm learned with Python 3rd: Radix conversion
Let's develop an investment algorithm with Python 1
Algorithm learned with Python 12th: Maze search
Excel with Python
Algorithm learned with Python 11th: Tree structure
Microcomputer with Python
Cast with python
Calculate the shortest route of a graph with Dijkstra's algorithm and Python
Algorithm learned with Python 13th: Tower of Hanoi
Algorithm learned with Python 16th: Sorting (insertion sort)
Algorithm learned with Python 14th: Tic-tac-toe (ox problem)
Algorithm learned with Python 15th: Sorting (selection sort)
Solving the Python knapsack problem with the greedy algorithm
Serial communication with Python
Zip, unzip with python
Django 1.11 started with Python3.6
Primality test with Python
Python with eclipse + PyDev.
Socket communication with Python
Data analysis with python 2
Scraping with Python (preparation)
A * algorithm (Python edition)
Python basic grammar / algorithm
Try scraping with Python.
Learning Python with ChemTHEATER 03
Sequential search with Python
"Object-oriented" learning with python
Run Python with VBA
Handling yaml with python
Solve AtCoder 167 with python
Serial communication with python
[Python] Use JSON with Python