Python program for Shortest path of a weighted graph where weight is 1 or 2

In this article, we are going to write code to find the shortest path of a weighted graph where weight is 1 or 2. since the weight is either 1 or 2. Whenever there is a weight of two, we will add an extra edge between them and make each weight to 1. Here we will first go through how to create a graph then we will use bfs and create the array of previously visited nodes. From the previously visited array, we will construct the path.

Graph Representation

So First we need to represent the graph in a way computationally feasible. I’m going to represent in an adjacency list. where for every node in the graph we will maintain a list of neighboring nodes. Let’s code

from collections import defaultdict
class Graph:

  def __init__(self,numVertices):
    self.numVertices = numVertices
    #For every node we had a list mapped to it with default dict.
    self.graph = defaultdict(list)

Here the graph variable contains a defaultdict with nodes mapping to list of neighboring edges. but we have to write a function to create edges and maintain lists for each. The below function will create that mapping. All the functions are written inside the Graph class.

def addEdge(self,edge_start,edge_end,edge_weight):

    if edge_weight==1:
      self.graph[edge_start].append(edge_end)
      #if edge weight is one the we directly add to list of neighbors of that particular node

    else:
      #if the edge weight is not one we will add one more edge and make both the edge weights to 1
      self.graph[edge_start].append(self.numVertices)
      # we are adding it as the last vertex and update it's neighbours
      self.graph[self.numVertices].append(edge_end)
      #we will increase the numVertices since one more vertex added
      self.numVertices+=1

Suppose there is an edge between node 4 and node 5 with weight 2. 
4--(2)-->5
we will make it as 4 --(1)--> 6 --(1)--> 5
so that all edges have weight 1.

The reason for changing the edge weights from 2 to 1  is we can make use of BFS to find the shortest path in a graph. If you don’t know the breadth-first search, Please go through this article first.

Shortest Path

During the breadth-first search we main an extra array to save the parent of each node, the index is the node, and value at index is the parent of the index. With the help of this array, we can construct the path. Let’s see the Python code:

def shortestPath(self,src,dest):
    visited = [False]*self.numVertices

    queue  = [src]
    visited[src]=True
    #prev is the extra array we maintain.
    prev = [None for i in range(self.numVertices)]

    while len(queue)!=0:
      s = queue.pop(0)

      for i in self.graph[s]:
        if visited[i]==False:
          queue.append(i)
          visited[i]=True
          # When we visited a node at that index we will have a value s 
          #since from s we are coming to this i
          prev[i] = s
          if i==dest:
            print(prev)
            #When we find the dest we will break 
            #and call construct path to get the path.
            print(self.ConstructPath(src,dest,prev))
            print("Found!!!")
            break

Now we have to construct the path from the extra array.

Construct Path

we will start with the index of destination and then we will go to the value of prev[index] as an index and continue till we find the source. while doing we will add to the path and we will reverse that to get the output. Let’s code:

def ConstructPath(self,src,dest,prev):
  path = [dest]
  index = prev[dest]
  path.append(index)
  count = len(prev)
  while(count>0):
    index = prev[index]
    path.append(index)
    count-=1
    if prev[index]==src:
      path.append(prev[index])
      path.reverse()
      return "-->".join(map(str,path))
  return "Not Found!"

So this is our way to solve this problem. Below is the overall code.

from collections import defaultdict
class Graph:

  def __init__(self,numVertices):
    self.numVertices = numVertices
    self.graph = defaultdict(list)

  def addEdge(self,edge_start,edge_end,edge_weight):
    if edge_weight==1:
      self.graph[edge_start].append(edge_end)
    else:
      self.graph[edge_start].append(self.numVertices)
      self.graph[self.numVertices].append(edge_end)
      self.numVertices+=1
  def printGraph(self):
    for i in range(self.numVertices):
      print(f"{i}--->{self.graph[i]} ")
  def shortestPath(self,src,dest):
    visited = [False]*self.numVertices

    queue  = [src]
    visited[src]=True
    prev = [None for i in range(self.numVertices)]

    while len(queue)!=0:
      s = queue.pop(0)

      for i in self.graph[s]:
        if visited[i]==False:
          queue.append(i)
          visited[i]=True
          prev[i] = s
          if i==dest:
            print(prev)
            print(self.ConstructPath(src,dest,prev))
            print("Found!!!")
            break
    print("Not Found!!")
  def ConstructPath(self,src,dest,prev):
    path = [dest]
    index = prev[dest]
    path.append(index)
    count = len(prev)
    while(count>0):
      index = prev[index]
      path.append(index)
      count-=1
      if prev[index]==src:
        path.append(prev[index])
        path.reverse()
        return "-->".join(map(str,path))
    return "Not Found!"

if __name__=='__main__':
  g = Graph(7)
  g.addEdge(0, 1, 1)
  g.addEdge(1,2,2)
  g.addEdge(1,3,1)
  g.addEdge(2,3,1)
  g.addEdge(3,6,1)
  g.addEdge(0,2,2)
  g.addEdge(0,5,1)
  g.addEdge(5,4,2)
  g.addEdge(4,3,2)
  g.addEdge(4,6,1) 
  g.printGraph()
  g.shortestPath(0,6)

Output:
0–>1–>3–>6
The input is the below graph:

Feel free to share your thoughts and doubts down in the comment section.