Shortest path algorithms

shortestpath.py

from heapq import heappop, heappush
from copy import deepcopy

Inf = float('inf')

def printmatrix(M):
  """
  Print a nicely formatted matrix.
  """
  for i in range(len(M)):
    for j in range(len(M)):
      if M[i][j] == Inf:
        print(u' \u221E', end=' ')        # Print nice infinity symbol ∞
      else:
        print('%2d' % M[i][j], end=' ')
    print()

def bellmanford(G, s):
  """
  Find single-source shortest paths in the graph G (adjacency matrix)
  using the Bellman-Ford algorithm.

  Negative edges allowed, but no negative cycles.
  """
  n = len(G)
  d = [Inf]*n                         # Estimated distances
  d[s] = 0                            # Source node has estimate 0
  p = {}                              # Predecessors
  for _ in range(n):                  # |V|-1 times
    for u in range(n):                # For all the edges
      for v in range(n):              # ...
        if d[v] > d[u] + G[u][v]:     # Relax
          d[v] = d[u] + G[u][v]
          p[v] = u                    # Update predecessor

  for u in range(n):
    for v in range(n):
      if d[u] + G[u][v] < d[v]:
        return False                  # Negative cycle

  return d

def dijkstra(G, s):
  """
  Find single-source shortest paths in the graph G (adjacency matrix)
  using Dijkstra's algorithm.

  No negative edges allowed.
  """
  n = len(G)
  Q = [(0, s)]                        # Priority Queue using binary heap
  p = {}                              # Predecessors
  dist = [Inf]*n                      # Estimated distances
  dist[s] = 0                         # Source node has estimate 0
  visited = set()                     # Visited vertices

  while Q:                            # While the queue is not empty
    _, u = heappop(Q)                 # Extract-Min
    visited.add(u)                    # Visited vertex u

    for v in range(n):                # For all adjacent vertices of u
      if v in visited:
        continue

      if dist[v] > dist[u] + G[u][v]: # Relax
        dist[v] = dist[u] + G[u][v]
        p[v] = u
      heappush(Q, (dist[v], v))       # Decrease-Key

  return dist                         # Distances from s to all vertices

def floydwarshall(G):
  """
  Find all-pairs shortest path in the graph G (adjacency matrix)
  using the Floyd-Warshall algorithm.

  Negative edges allowed, but no negative cycles.
  """
  W = deepcopy(G)
  n = len(W)
  for k in range(n):
    for i in range(n):
      for j in range(n):
        W[i][j] = min(W[i][j], W[i][k] + W[k][j])
  return W

if __name__ == '__main__':
  # Figure 25.1 in Introduction to Algorithms 3e
  G = [
    [0,     3,   8, Inf,  -4],
    [Inf,   0, Inf,   1,   7],
    [Inf,   4,   0, Inf, Inf],
    [2,   Inf,  -5,   0, Inf],
    [Inf, Inf, Inf,   6,   0]
  ]

  printmatrix(floydwarshall(G))
  print(dijkstra(G, 0))
  print(bellmanford(G, 0))