Dijkstra algorithm: How to implement it with Python (solved with all explanations) ?

In this post, I will show you how to implement Dijkstra's algorithm for shortest path calculations in a graph with Python.

NB: If you need to revise how Dijstra's work, have a look to the post where I detail Dijkstra's algorithm operations step by step on the whiteboard, for the example below.

First, let's choose the right data structures. The implementation can be a nightmare if you pick a bad data structure for describing the graphs. In python, the dict structure is particularly handy, so let's use it.



On the above example, s has two neighbors a and b, linked with cost 2 and 1 respectively. So, our graph dictionary will look like this:

graph = {'s': {'a': 2, 'b': 1},
            'a': {'s': 3, 'b': 4, 'c':8},
            'b': {'s': 4, 'a': 2, 'd': 2},
            'c': {'a': 2, 'd': 7, 't': 4},
            'd': {'b': 1, 'c': 11, 't': 5},
            't': {'c': 3, 'd': 5}}

Now, if you have followed my previous post detailing the pseudo code and operations of Dijsktra's shortest path tree algorithm, you know that we need to keep track of  the following information:
- which nodes are not visited yet: we can use a list for that.
- a predecessor and a cost for every node. Let's use dict type for these elements, it will ease our life when we look up their value for a given node.

So the signature of our function is the following:

def shortestpath(graph,src,dest,visited=[],distances={},predecessors={}):

It will return the shortest path from src to dest, but could as well return a shortest path tree.

You will need to know the two following python functions to implement Dijkstra smartly.

mydict.get(mykeyvalue,17)

this function of a dict element (here 'mydict') searches for the value of the dict for the keyvalue 'mykeyvalue'. If this key does not exist in the dict, the function does not raise an error. Instead, it returns the value provided as the second argument, here 17.

min(mydict, key=mydict.get)

The min function can be applied on dictionaries! You just have to specify the key function as in the above example.

Now, we have everything we need to implement the function. I propose that you try it on your side and then look at the solution. Otherwise, you'll not memorize the algorithm. That being said, it's your choice...

One last thing before jumping to the code: if your graph bears negative weights, Dijkstra's algorithm will run into trouble - you should rather use Floyd or Bellman-Ford algorithm, which are able to detect negative cycles in graphs. Remember that Dijsktra's algorithm is for graphs with positive weights only.

def dijkstra(graph,src,dest,visited=[],distances={},predecessors={}):
    """ calculates a shortest path tree routed in src
    """    
    # a few sanity checks
    if src not in graph:
        raise TypeError('The root of the shortest path tree cannot be found')
    if dest not in graph:
        raise TypeError('The target of the shortest path cannot be found')    
    # ending condition
    if src == dest:
        # We build the shortest path and display it
        path=[]
        pred=dest
        while pred != None:
            path.append(pred)
            pred=predecessors.get(pred,None)
        print('shortest path: '+str(path)+" cost="+str(distances[dest])) 
    else :     
        # if it is the initial  run, initializes the cost
        if not visited: 
            distances[src]=0
        # visit the neighbors
        for neighbor in graph[src] :
            if neighbor not in visited:
                new_distance = distances[src] + graph[src][neighbor]
                if new_distance < distances.get(neighbor,float('inf')):
                    distances[neighbor] = new_distance
                    predecessors[neighbor] = src
        # mark as visited
        visited.append(src)
        # now that all neighbors have been visited: recurse                         
        # select the non visited node with lowest distance 'x'
        # run Dijskstra with src='x'
        unvisited={}
        for k in graph:
            if k not in visited:
                unvisited[k] = distances.get(k,float('inf'))        
        x=min(unvisited, key=unvisited.get)
        dijkstra(graph,x,dest,visited,distances,predecessors)
        


if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    #unittest.main()
    graph = {'s': {'a': 2, 'b': 1},
            'a': {'s': 3, 'b': 4, 'c':8},
            'b': {'s': 4, 'a': 2, 'd': 2},
            'c': {'a': 2, 'd': 7, 't': 4},
            'd': {'b': 1, 'c': 11, 't': 5},
            't': {'c': 3, 'd': 5}}
    dijkstra(graph,'s','t')

And magic... It returns:

shortest path: ['t', 'd', 'b', 's'] cost=8

Well done and congratulations for reading this post to the end! Now, if you are working on Dijsktra's algorithm, I think you are preparing for an interview at Google, Microsoft, or Amazon. Am I right?

Then, the best tip I can give you is to look at two books "Cracking the PM Interview: How to Land a Product Manager Job in Technology" and "Cracking the Coding Interview: 150 Programming Questions and Solutions". They were written by Gayle Laakmann McDowell, a former recruiter from Google who also worked at Apple and I find them really great!