from math import inf from itertools import product def floyd_warshall(n, edge): rn = range(n) dist = [[inf] * n for i in rn] nxt = [[0] * n for i in rn] for i in rn: dist[i][i] = 0 for u, v, w in edge: dist[u-1][v-1] = w nxt[u-1][v-1] = v-1 for k, i, j in product(rn, repeat=3): sum_ik_kj = dist[i][k] + dist[k][j] if dist[i][j] > sum_ik_kj: dist[i][j] = sum_ik_kj nxt[i][j] = nxt[i][k] print("pair dist path") for i, j in product(rn, repeat=2): if i != j: path = [i] while path[-1] != j: path.append(nxt[path[-1]][j]) print("%d → %d %4d %s" % (i + 1, j + 1, dist[i][j], ' → '.join(str(p + 1) for p in path))) if __name__ == '__main__': floyd_warshall(4, [[1, 3, -2], [2, 1, 4], [2, 3, 3], [3, 4, 2], [4, 2, -1]])