Files
troupe-standup/wizard-work/beam/bellman_ford.py
T

85 lines
2.7 KiB
Python

"""Bellman-Ford shortest paths with negative-cycle detection.
Given a directed weighted graph (possibly with negative edge weights) and a
source vertex, compute shortest-path distances from the source. Detects
negative-weight cycles reachable from the source.
Returns (dist, parent, has_negative_cycle). dist maps vertex -> distance (inf
if unreachable). parent maps vertex -> predecessor on a shortest path (or
None for source/unreachable). has_negative_cycle is True iff a negative cycle
reachable from the source exists.
"""
from __future__ import annotations
from typing import Dict, Hashable, List, Optional, Tuple
import math
Vertex = Hashable
Edge = Tuple[Vertex, Vertex, float] # (u, v, weight)
def bellman_ford(edges: List[Edge], n: int, source: Vertex,
vertices: Optional[List[Vertex]] = None) -> Tuple[Dict, Dict, bool]:
if vertices is None:
# infer vertices 0..n-1 unless source not in range
vertices = list(range(n))
INF = float("inf")
dist: Dict[Vertex, float] = {v: INF for v in vertices}
parent: Dict[Vertex, Optional[Vertex]] = {v: None for v in vertices}
if source not in dist:
dist[source] = 0.0
parent[source] = None
if source not in vertices:
vertices = vertices + [source]
# ensure dist/parent cover source
if source not in dist:
dist[source] = 0.0
parent[source] = None
dist[source] = 0.0
# Relax all edges V-1 times.
V = len(vertices)
for _ in range(V - 1):
updated = False
for u, v, w in edges:
if dist.get(u, INF) != INF and dist[u] + w < dist.get(v, INF):
dist[v] = dist[u] + w
parent[v] = u
updated = True
if not updated:
break
# Detect negative cycle reachable from source.
has_neg = False
for u, v, w in edges:
if dist.get(u, INF) != INF and dist[u] + w < dist.get(v, INF):
has_neg = True
break
return dist, parent, has_neg
def reconstruct_path(parent: Dict, source: Vertex, target: Vertex) -> Optional[List[Vertex]]:
if target not in parent:
return None
# walk
path = []
cur: Optional[Vertex] = target
seen = set()
while cur is not None:
if cur in seen:
return None # cycle
seen.add(cur)
path.append(cur)
if cur == source:
break
cur = parent.get(cur)
if cur is None and path[-1] != source:
# unreachable
if target != source and parent.get(target) is None and target != source:
return None
break
if path and path[-1] != source:
return None
path.reverse()
return path