Skip to content

Instantly share code, notes, and snippets.

@kachayev
Last active November 29, 2023 10:21
Star You must be signed in to star a gist
Embed
What would you like to do?
Dijkstra shortest path algorithm based on python heapq heap implementation
from collections import defaultdict
from heapq import *
def dijkstra(edges, f, t):
g = defaultdict(list)
for l,r,c in edges:
g[l].append((c,r))
q, seen, mins = [(0,f,())], set(), {f: 0}
while q:
(cost,v1,path) = heappop(q)
if v1 not in seen:
seen.add(v1)
path = (v1, path)
if v1 == t: return (cost, path)
for c, v2 in g.get(v1, ()):
if v2 in seen: continue
prev = mins.get(v2, None)
next = cost + c
if prev is None or next < prev:
mins[v2] = next
heappush(q, (next, v2, path))
return float("inf"), None
if __name__ == "__main__":
edges = [
("A", "B", 7),
("A", "D", 5),
("B", "C", 8),
("B", "D", 9),
("B", "E", 7),
("C", "E", 5),
("D", "E", 15),
("D", "F", 6),
("E", "F", 8),
("E", "G", 9),
("F", "G", 11)
]
print "=== Dijkstra ==="
print edges
print "A -> E:"
print dijkstra(edges, "A", "E")
print "F -> G:"
print dijkstra(edges, "F", "G")
@andreppires
Copy link

andreppires commented Feb 21, 2018

If you want a clean output you can do it by adding this lines:

out = dijkstra(edges, start, end)
data = {}
data['cost']=out[0]
aux=[]
while len(out)>1:
    aux.append(out[0])
    out = out[1]
aux.remove(data['cost'])
aux.reverse()
data['path']=aux

@romech
Copy link

romech commented Apr 2, 2018

@andreppires, here is a shorter version

make_path = lambda tup: (*make_path(tup[1]), tup[0]) if tup else ()
out = dijkstra(edges, start, end)
path = make_path(out[1])

@GGonzalezRojas
Copy link

Very good algorithm, it helped me a lot in a task. I made a translation commenting on the Spanish code for a better understanding.

@YDD9
Copy link

YDD9 commented May 5, 2018

@hhu94 line 12 is not redundant either. For an existing node in q, heappush will keep adding different costs for that node, so without line 12, that node will be visited again and update with a higher cost later.

@amandalmia14
Copy link

amandalmia14 commented Jun 2, 2018

Can it be possible to optimise more? As I am getting run-time error(NZEC) in codechef.
I am working on this https://www.codechef.com/INOIPRAC/

@whiledoing
Copy link

@licensed @andreppires @romech

It may very simple by change line 14 into path += (v1, ), this will make output more clear and reverse the path in the meanwhile.

A -> E:
(14, ('A', 'B', 'E'))
F -> G:
(11, ('F', 'G'))

@whiledoing
Copy link

@waylonflinn

I think you are right. The key problem here is when node v2 is already in the heap, you should not put v2 into heap again, instead you need to heap.remove(v) and then head.insert(v2) if new cost of v2 is better then original cost of v2 recorded in the heap. But indeed remove node in heap is just O(n), so that will not be any better then original implementation of Dijkstra using distance array.

I change the code by taking the distance array into consideration which will record the min value of each node already put into the heap. So we will only put the v2 into the heap just on new value is better then the original one which I think in most case it will improve the performance of this algorithm. But just as @tjwudi mentioned, in worst case, it still will be O(V^2 logV) :)

Revised version:

def dijkstra_revised(edges, f, t):
    g = defaultdict(list)
    for l,r,c in edges:
        g[l].append((c,r))

    # dist records the min value of each node in heap.
    q, seen, dist = [(0,f,())], set(), {f: 0}
    while q:
        (cost,v1,path) = heappop(q)
        if v1 in seen: continue

        seen.add(v1)
        path += (v1,)
        if v1 == t: return (cost, path)

        for c, v2 in g.get(v1, ()):
            if v2 in seen: continue

            # Not every edge will be calculated. The edge which can improve the value of node in heap will be useful.
            if v2 not in dist or cost+c < dist[v2]:
                dist[v2] = cost+c
                heappush(q, (cost+c, v2, path))

    return float("inf")

@kachayev
Copy link
Author

@whiledoing Thanks! I've made an adjustment to the initial gist (slightly changed to avoid checking the same key from dist twice).

@VeNoMouS
Copy link

I dunno, @whiledoing's tuple output doesn't look as defunct as your latest change @kachayev.... I started getting some weird nested tuples with your version.

@kevthanewversi
Copy link

this is brilliant. Thanks @kachayev

@u2takey
Copy link

u2takey commented Nov 13, 2018

seen is redundant

def dijkstra(edges, f, t):
    g = defaultdict(list)
    for l,r,c in edges:
        g[l].append((c,r))
  
    q, dist = [(0,f,())], {f: 0}
    while q:
        (cost, node, path) = heappop(q) 
        if cost > dist[node]: continue
        path += (node, )
        if node == t: 
            return (cost, path)
        for w, n in g.get(node, ()):
            oldc = dist.get(n, float("inf"))
            newc = cost + w
            if newc < oldc:
                dist[n] = newc # relax
                heappush(q, (newc, n, path))

    return float("inf")

or you can just use seen, ignore mins/dist

from collections import defaultdict
from heapq import *

def dijkstra(edges, f, t):
    g = defaultdict(list)
    for l,r,c in edges:
        g[l].append((c,r))

    q, seen = [(0,f,())], set()
    while q:
        (cost,v1,path) = heappop(q)
        if v1 not in seen:
            seen.add(v1)
            path = (v1, path)
            if v1 == t: return (cost, path)

            for c, v2 in g.get(v1, ()):
                if v2 in seen: continue
                next = cost + c
                heappush(q, (next, v2, path))

    return float("inf")

@lycutter
Copy link

Thanks for your code very much. But I want to make some expansion on this basis. The situation is that our map is a matrix, and there are more than one shortest path to reach the destination, if I want to find all the road not just the one, how to modify the code to achieve this? Thanks again.

@nolfonzo
Copy link

This is a slightly simpler approach, following the wikipedia definition closely:
http://rebrained.com/?p=392

import sys
def shortestpath(graph,start,end,visited=[],distances={},predecessors={}):
"""Find the shortest path btw start & end nodes in a graph"""

# detect if first time through, set current distance to zero
if not visited: distances[start]=0

# if we've found our end node, find the path to it, and return
if start==end:
    path=[]
    while end != None:
        path.append(end)
        end=predecessors.get(end,None)
    return distances[start], path[::-1]

# process neighbors as per algorithm, keep track of predecessors
for neighbor in graph[start]:
    if neighbor not in visited:
        neighbordist = distances.get(neighbor,sys.maxint)
        tentativedist = distances[start] + graph[start][neighbor]
        if tentativedist < neighbordist:
            distances[neighbor] = tentativedist
            predecessors[neighbor]=start

# neighbors processed, now mark the current node as visited 
visited.append(start)

# finds the closest unvisited node to the start 
unvisiteds = dict((k, distances.get(k,sys.maxint)) for k in graph if k not in visited)
closestnode = min(unvisiteds, key=unvisiteds.get)

# now take the closest node and recurse, making it current 
return shortestpath(graph,closestnode,end,visited,distances,predecessors)

if name == "main":
graph = {'a': {'w': 14, 'x': 7, 'y': 9},
'b': {'w': 9, 'z': 6},
'w': {'a': 14, 'b': 9, 'y': 2},
'x': {'a': 7, 'y': 10, 'z': 15},
'y': {'a': 9, 'w': 2, 'x': 10, 'z': 11},
'z': {'b': 6, 'x': 15, 'y': 11}}
print shortestpath(graph,'a','a')
print shortestpath(graph,'a','b')

"""
Expected Result:
    (0, ['a']) 
    (20, ['a', 'y', 'w', 'b'])
    """

@daidai21
Copy link

Flatten path version:

from collections import defaultdict
from heapq import *


def dijkstra(edges, f, t):
    g = defaultdict(list)
    for l, r, c in edges:
        g[l].append((c, r))

    q, seen, mins = [(0, f, [])], set(), {f: 0}
    while q:
        (cost, v1, path) = heappop(q)
        if v1 not in seen:
            seen.add(v1)
            path = [v1] + path
            if v1 == t:
                return (cost, path)

            for c, v2 in g.get(v1, ()):
                if v2 in seen:
                    continue
                prev = mins.get(v2, None)
                next = cost + c
                if prev is None or next < prev:
                    mins[v2] = next
                    heappush(q, (next, v2, path))

    return (float("inf"), [])


if __name__ == "__main__":
    edges = [
        ("A", "B", 7),
        ("A", "D", 5),
        ("B", "C", 8),
        ("B", "D", 9),
        ("B", "E", 7),
        ("C", "E", 5),
        ("D", "E", 15),
        ("D", "F", 6),
        ("E", "F", 8),
        ("E", "G", 9),
        ("F", "G", 11)
    ]

    print("=== Dijkstra ===")
    print(edges)
    print("A -> E: ", end="")
    print(dijkstra(edges, "A", "E"))
    print("F -> G: ", end="")
    print(dijkstra(edges, "F", "G"))
=== Dijkstra ===
[('A', 'B', 7), ('A', 'D', 5), ('B', 'C', 8), ('B', 'D', 9), ('B', 'E', 7), ('C', 'E', 5), ('D', 'E', 15), ('D', 'F', 6), ('E', 'F', 8), ('E', 'G', 9), ('F', 'G', 11)]
A -> E: (14, ['E', 'B', 'A'])
F -> G: (11, ['G', 'F'])

@ehborisov
Copy link

ehborisov commented Feb 28, 2020

Hi, I think I made a bit cleaner (subjectively :)) implementation in Python that uses RBTree as a priority queue with tests there

https://github.com/ehborisov/algorithms/blob/master/8.Graphs/dijkstra.py

@coldmanck
Copy link

coldmanck commented Mar 6, 2020

Unless I am missing something here, this is a BFS with a min-heap, not a Dijkstra's algorithm.

@JixinSiND Dijkstra's algorithm is essentially a weighted version of BFS.

@alelom
Copy link

alelom commented Jun 6, 2020

Just leaving a comment to let the author know that his code has been inappropriately taken and re-used as material for teaching at a University master in London. The authorship has been modified to report the lecturer's one instead.
https://www.dcs.bbk.ac.uk/~ale/pwd/2019-20/pwd-8/src/pwd-ex-dijkstra+heap.py

@kachayev
Copy link
Author

kachayev commented Jun 8, 2020

@alelom Thanks a lot for letting me know, such a kind of you! This is not the first time this code was copy-pasted into lecture materials and/or projects codebases. Honestly, if it helped students to learn - I would be glad and proud. I care less about authorship or any sort of attribution. On the one hand, I wouldn't want to encourage disrespectful actions, on the other hand, I don't have reliable way to prevent this from happening. So, choosing between spread of knowledge or nurturing morality, I would always vote for the former. Thanks again for letting me know!

@alelom
Copy link

alelom commented Jun 9, 2020

I just care for what is right. If I were the lecturer, I'd quote the real author and the source – an action that does not diminish the teaching potential, and encourages sharing of good code lawfully.

@Andy9822
Copy link

Thank you so much for this gift, very clean and clever solution 😄

If anyone just wonders how to easily receive as output only the value of the solution remove the cost from the return at line 15:

if v1 == t: return cost
instead of
if v1 == t: return (cost, path)

@konmaz
Copy link

konmaz commented Nov 6, 2020

Nice and clean

@PoTseCheng
Copy link

Thank you very much for this beautiful algorithm.

@xdavidliu
Copy link

pretty sure this is not Dijkstra; you're doing heappush(q, (next, v2, path)) at the very end, but in True dijkstra it would need a call to "decrease_key", which in python is heap._siftdown

@chausen
Copy link

chausen commented Feb 27, 2022

@xdavidliu I was confused by this until I saw https://stackoverflow.com/a/31123108. I think Dijkstra's algorithm is a higher level concept, so either implementation is valid.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment