use alloc::collections::VecDeque;
use super::{bellman_ford::Paths, BoundedMeasure, NegativeCycle};
use crate::prelude::*;
use crate::visit::{IntoEdges, IntoNodeIdentifiers, NodeIndexable};
use alloc::{vec, vec::Vec};
pub fn spfa<G, F, K>(
graph: G,
source: G::NodeId,
edge_cost: F,
) -> Result<Paths<G::NodeId, K>, NegativeCycle>
where
G: IntoEdges + IntoNodeIdentifiers + NodeIndexable,
F: FnMut(G::EdgeRef) -> K,
K: BoundedMeasure + Copy,
{
let ix = |i| graph.to_index(i);
let pred = vec![None; graph.node_bound()];
let mut dist = vec![K::max(); graph.node_bound()];
dist[ix(source)] = K::default();
let mut queue: VecDeque<G::NodeId> = VecDeque::with_capacity(graph.node_bound());
let mut in_queue = vec![false; graph.node_bound()];
queue.push_back(source);
in_queue[ix(source)] = true;
let (distances, predecessors) = spfa_loop(graph, dist, Some(pred), queue, in_queue, edge_cost)?;
Ok(Paths {
distances,
predecessors: predecessors.unwrap_or_default(),
})
}
#[allow(clippy::type_complexity)]
pub(crate) fn spfa_loop<G, F, K>(
graph: G,
mut distances: Vec<K>,
mut predecessors: Option<Vec<Option<G::NodeId>>>,
mut queue: VecDeque<G::NodeId>,
mut in_queue: Vec<bool>,
mut edge_cost: F,
) -> Result<(Vec<K>, Option<Vec<Option<G::NodeId>>>), NegativeCycle>
where
G: IntoEdges + IntoNodeIdentifiers + NodeIndexable,
F: FnMut(G::EdgeRef) -> K,
K: BoundedMeasure + Copy,
{
let ix = |i| graph.to_index(i);
let mut visits = vec![0; graph.node_bound()];
while let Some(i) = queue.pop_front() {
in_queue[ix(i)] = false;
if visits[ix(i)] >= graph.node_bound() {
return Err(NegativeCycle(()));
}
visits[ix(i)] += 1;
for edge in graph.edges(i) {
let j = edge.target();
let w = edge_cost(edge);
let (dist, overflow) = distances[ix(i)].overflowing_add(w);
if !overflow && dist < distances[ix(j)] {
distances[ix(j)] = dist;
if let Some(p) = predecessors.as_mut() {
p[ix(j)] = Some(i)
}
if !in_queue[ix(j)] {
in_queue[ix(j)] = true;
queue.push_back(j);
}
}
}
}
Ok((distances, predecessors))
}