Floyd-Marshall algorithm for all-pairs shortest paths

floyd_marshall.py

import numpy as np
import numba

@numba.jit
def floyd_marshall(graph):
    n = len(graph.nodes)
    A = np.full((n, n, n), fill_value=np.inf)
    
    # base cases
    for edge in graph.edges:
        A[edge.node1, edge.node2, 0] = edge.cost            
    np.fill_diagonal(A[:, :, 0], 0)
                
    for k in range(1, n):
        for i in range(n):
            for j in range(n):
                A[i, j, k] = min(A[i, j, k-1], A[i, k, k-1] + A[k, j, k-1])
                
    if (np.diag(A[:, :, -1]) == 0).all():
        return A[:, :, -1]
    
    else:
        raise ValueError('Negative cycle present')