55 lines
1.3 KiB
Python
55 lines
1.3 KiB
Python
"""Floyd-Warshall all-pairs shortest paths."""
|
|
from __future__ import annotations
|
|
|
|
|
|
def floyd_warshall(n, edges):
|
|
INF = float("inf")
|
|
dist = [[INF] * n for _ in range(n)]
|
|
nxt = [[None] * n for _ in range(n)]
|
|
for i in range(n):
|
|
dist[i][i] = 0
|
|
nxt[i][i] = i
|
|
for (u, v, w) in edges:
|
|
if w < dist[u][v]:
|
|
dist[u][v] = w
|
|
nxt[u][v] = v
|
|
for k in range(n):
|
|
dk = dist[k]
|
|
nk = nxt[k]
|
|
for i in range(n):
|
|
dik = dist[i][k]
|
|
if dik == INF:
|
|
continue
|
|
di = dist[i]
|
|
ni = nxt[i]
|
|
for j in range(n):
|
|
if dk[j] == INF:
|
|
continue
|
|
nd = dik + dk[j]
|
|
if nd < di[j]:
|
|
di[j] = nd
|
|
ni[j] = ni[k]
|
|
for i in range(n):
|
|
if dist[i][i] < 0:
|
|
raise ValueError("negative cycle detected")
|
|
return dist, nxt
|
|
|
|
|
|
def reconstruct_path(nxt, s, t):
|
|
if nxt[s][t] is None:
|
|
return None
|
|
path = [s]
|
|
cur = s
|
|
seen = set()
|
|
while cur != t:
|
|
if cur in seen:
|
|
return None
|
|
seen.add(cur)
|
|
cur = nxt[cur][t]
|
|
if cur is None:
|
|
return None
|
|
path.append(cur)
|
|
if len(path) > len(nxt) + 1:
|
|
return None
|
|
return path
|