From e98f893f1513391dc43e4787c5d79bee7f4069a4 Mon Sep 17 00:00:00 2001 From: tachsin Date: Fri, 14 Aug 2026 00:06:40 +0300 Subject: [PATCH] feat: add BMSSP single-source shortest paths Adds sssp / sssp_all following Duan et al. (arXiv:2504.17033), with the same successor-function API as Dijkstra. Distances match Dijkstra on finite reachable graphs. Co-authored-by: Cursor --- src/directed/mod.rs | 1 + src/directed/sssp.rs | 529 +++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 3 + tests/sssp.rs | 105 +++++++++ 4 files changed, 638 insertions(+) create mode 100644 src/directed/sssp.rs create mode 100644 tests/sssp.rs diff --git a/src/directed/mod.rs b/src/directed/mod.rs index 36086a60..d6359478 100644 --- a/src/directed/mod.rs +++ b/src/directed/mod.rs @@ -13,6 +13,7 @@ pub mod edmonds_karp; pub mod fringe; pub mod idastar; pub mod iddfs; +pub mod sssp; pub mod strongly_connected_components; pub mod topological_sort; pub mod yen; diff --git a/src/directed/sssp.rs b/src/directed/sssp.rs new file mode 100644 index 00000000..873f3ac4 --- /dev/null +++ b/src/directed/sssp.rs @@ -0,0 +1,529 @@ +//! Single-source shortest paths via the BMSSP algorithm of Duan, Mao, Mao, +//! Shu and Yin ([arXiv:2504.17033](https://arxiv.org/abs/2504.17033)). +//! +//! The public functions have the same shape as [`dijkstra_all`](super::dijkstra::dijkstra_all) +//! and [`dijkstra`](super::dijkstra::dijkstra). Distances are exact. The internal +//! frontier queue is a balanced tree rather than the block structure of the +//! paper, so this implementation does not claim the \(O(m\log^{2/3}n)\) bound. +//! A linear repair pass fixes nodes left stale when many paths share a length +//! (the paper assumes unique path lengths). + +use crate::FxIndexMap; +use indexmap::map::Entry::{Occupied, Vacant}; +use num_traits::Zero; +use rustc_hash::{FxHashMap, FxHashSet}; +use std::cmp::Reverse; +use std::collections::{BTreeMap, BinaryHeap, HashMap, VecDeque}; +use std::hash::Hash; + +/// Compute a shortest path using the [BMSSP](https://arxiv.org/abs/2504.17033) +/// single-source algorithm. +/// +/// Same result as [`dijkstra`](super::dijkstra::dijkstra) on a **finite** +/// reachable graph: the path from `start` to a cheapest node for which +/// `success` is true, together with its cost. +/// +/// This computes all distances first. Prefer +/// [`dijkstra`](super::dijkstra::dijkstra) when the successor graph is +/// unbounded or you only need one target. +/// +/// # Example +/// +/// ``` +/// use pathfinding::prelude::sssp; +/// +/// fn successors(&n: &u32) -> Vec<(u32, usize)> { +/// if n <= 4 { +/// vec![(n * 2, 10), (n * 2 + 1, 10)] +/// } else { +/// vec![] +/// } +/// } +/// +/// let (path, cost) = sssp(&1, successors, |&n| n == 9).expect("no path"); +/// assert_eq!(cost, 30); +/// assert_eq!(path, vec![1, 2, 4, 9]); +/// ``` +pub fn sssp(start: &N, successors: FN, mut success: FS) -> Option<(Vec, C)> +where + N: Eq + Hash + Clone, + C: Zero + Ord + Copy, + FN: FnMut(&N) -> IN, + IN: IntoIterator, + FS: FnMut(&N) -> bool, +{ + if success(start) { + return Some((vec![start.clone()], Zero::zero())); + } + let parents = sssp_all(start, successors); + let (target, cost) = parents + .iter() + .filter(|(node, _)| success(node)) + .min_by_key(|(_, (_, cost))| *cost) + .map(|(node, (_, cost))| (node.clone(), *cost))?; + Some((super::dijkstra::build_path(&target, &parents), cost)) +} + +/// Determine all reachable nodes from a starting point, and an optimal parent +/// and cost for each, using BMSSP ([arXiv:2504.17033](https://arxiv.org/abs/2504.17033)). +/// +/// Same result type as [`dijkstra_all`](super::dijkstra::dijkstra_all): every +/// reachable node except `start` maps to `(parent, cost_from_start)`. Use +/// [`build_path`](super::dijkstra::build_path) to recover a path. +/// +/// The reachable graph must be finite. +/// +/// # Example +/// +/// ``` +/// use pathfinding::prelude::sssp_all; +/// +/// fn successors(&n: &u32) -> Vec<(u32, usize)> { +/// if n <= 4 { +/// vec![(n * 2, 10), (n * 2 + 1, 10)] +/// } else { +/// vec![] +/// } +/// } +/// +/// let reachables = sssp_all(&1, successors); +/// assert_eq!(reachables.len(), 8); +/// assert_eq!(reachables[&2], (1, 10)); +/// assert_eq!(reachables[&9], (4, 30)); +/// ``` +pub fn sssp_all(start: &N, successors: FN) -> HashMap +where + N: Eq + Hash + Clone, + C: Zero + Ord + Copy, + FN: FnMut(&N) -> IN, + IN: IntoIterator, +{ + let (nodes, adj) = materialize(start, successors); + let n = nodes.len(); + if n == 0 { + return HashMap::new(); + } + + let mut dist = vec![None; n]; + let mut pred = vec![None; n]; + dist[0] = Some(Zero::zero()); + + let (k, t, lmax) = parameters(n); + let mut engine = Engine { + adj: &adj, + dist: &mut dist, + pred: &mut pred, + k, + t, + }; + engine.bmssp(lmax, None, &[0]); + engine.repair(); + + let mut out = HashMap::with_capacity(n.saturating_sub(1)); + for i in 1..n { + if let (Some(p), Some(cost)) = (pred[i], dist[i]) { + out.insert(nodes[i].clone(), (nodes[p].clone(), cost)); + } + } + out +} + +#[expect( + clippy::cast_possible_truncation, + clippy::cast_precision_loss, + clippy::cast_sign_loss +)] +fn parameters(n: usize) -> (usize, usize, usize) { + let n = n.max(2) as f64; + let log_n = n.log2(); + let k = log_n.powf(1.0 / 3.0).floor() as usize; + let t = log_n.powf(2.0 / 3.0).floor() as usize; + let k = k.max(1); + let t = t.max(1); + let lmax = (log_n / t as f64).ceil() as usize; + (k, t, lmax.max(1)) +} + +fn materialize(start: &N, mut successors: FN) -> (Vec, Vec>) +where + N: Eq + Hash + Clone, + C: Copy, + FN: FnMut(&N) -> IN, + IN: IntoIterator, +{ + let mut nodes: FxIndexMap = FxIndexMap::default(); + nodes.insert(start.clone(), ()); + let mut adj = Vec::new(); + let mut i = 0; + while i < nodes.len() { + let node = nodes.get_index(i).unwrap().0.clone(); + let mut edges = Vec::new(); + for (succ, cost) in successors(&node) { + let j = match nodes.entry(succ) { + Vacant(e) => { + let idx = e.index(); + e.insert(()); + idx + } + Occupied(e) => e.index(), + }; + edges.push((j, cost)); + } + adj.push(edges); + i += 1; + } + (nodes.into_iter().map(|(n, ())| n).collect(), adj) +} + +fn exp2_cap(exp: usize, cap: usize) -> usize { + if exp >= usize::BITS as usize { + cap + } else { + (1usize << exp).min(cap) + } +} + +fn less_than(value: &C, bound: Option) -> bool { + match bound { + None => true, + Some(b) => *value < b, + } +} + +struct Engine<'a, C> { + adj: &'a [Vec<(usize, C)>], + dist: &'a mut [Option], + pred: &'a mut [Option], + k: usize, + t: usize, +} + +impl Engine<'_, C> +where + C: Zero + Ord + Copy, +{ + fn relax(&mut self, u: usize, v: usize, weight: C) -> Option { + let du = self.dist[u]?; + let cand = du + weight; + let valid = self.dist[v].is_none_or(|old| cand <= old); + if !valid { + return None; + } + let better = self.dist[v] + .is_none_or(|old| cand < old || (cand == old && self.pred[v].is_none_or(|p| u < p))); + if better { + self.dist[v] = Some(cand); + self.pred[v] = Some(u); + } + Some(cand) + } + + fn bmssp( + &mut self, + level: usize, + bound: Option, + sources: &[usize], + ) -> (Option, Vec) { + if sources.is_empty() { + return (bound, Vec::new()); + } + if level == 0 { + return self.base_case(bound, sources[0]); + } + + let (mut pivots, witnessed) = self.find_pivots(bound, sources); + if pivots.is_empty() { + pivots.extend(sources.iter().copied()); + } + let m = exp2_cap((level - 1).saturating_mul(self.t), self.adj.len().max(1)); + let mut queue = PartialQueue::new(bound, m.max(1)); + for &x in &pivots { + if let Some(dx) = self.dist[x] { + queue.insert(x, dx); + } + } + + let limit = self.k.saturating_mul(exp2_cap( + level.saturating_mul(self.t), + self.adj.len().max(1), + )); + let mut completed = FxHashSet::default(); + let mut last_prime = bound; + + while completed.len() < limit && !queue.is_empty() { + let (si, bi) = queue.pull(); + if si.is_empty() { + break; + } + let (bi_prime, ui) = self.bmssp(level - 1, bi, &si); + last_prime = bi_prime; + completed.extend(ui.iter().copied()); + + let mut batch = Vec::new(); + for &u in &ui { + for &(v, weight) in &self.adj[u] { + let Some(cand) = self.relax(u, v, weight) else { + continue; + }; + if !less_than(&cand, bound) { + continue; + } + if less_than(&cand, bi) { + batch.push((v, cand)); + } else { + queue.insert(v, cand); + } + } + } + for &x in &si { + if let Some(dx) = + self.dist[x].filter(|dx| !less_than(dx, bi_prime) && less_than(dx, bi)) + { + batch.push((x, dx)); + } + } + queue.batch_prepend(&batch); + } + + let b_prime = match (last_prime, bound) { + (Some(a), Some(b)) => Some(a.min(b)), + (None, b) | (b, None) => b, + }; + for &x in &witnessed { + if self.dist[x].is_some_and(|dx| less_than(&dx, b_prime)) { + completed.insert(x); + } + } + (b_prime, completed.into_iter().collect()) + } + + fn base_case(&mut self, bound: Option, source: usize) -> (Option, Vec) { + let mut seen = FxHashSet::default(); + seen.insert(source); + let mut heap = BinaryHeap::new(); + if let Some(ds) = self.dist[source] { + heap.push(Reverse((ds, source))); + } + + while seen.len() < self.k + 1 { + let Some(Reverse((cost, u))) = heap.pop() else { + break; + }; + let Some(du) = self.dist[u] else { + continue; + }; + if cost > du || !less_than(&du, bound) { + continue; + } + seen.insert(u); + for &(v, weight) in &self.adj[u] { + let Some(cand) = self.relax(u, v, weight) else { + continue; + }; + if less_than(&cand, bound) { + heap.push(Reverse((cand, v))); + } + } + } + + if seen.len() <= self.k { + (bound, seen.into_iter().collect()) + } else { + let b_prime = seen.iter().filter_map(|&v| self.dist[v]).max(); + let u: Vec = seen + .into_iter() + .filter(|&v| self.dist[v].is_some_and(|dv| less_than(&dv, b_prime))) + .collect(); + (b_prime, u) + } + } + + /// Propagate leftover improvements. BMSSP can leave a child stale when a + /// parent is later corrected; that happens on graphs with many equal-cost + /// paths, which the paper excludes by assuming unique path lengths. + fn repair(&mut self) { + let n = self.adj.len(); + let mut queue = VecDeque::new(); + let mut queued = vec![false; n]; + for (u, queued_flag) in queued.iter_mut().enumerate() { + if self.dist[u].is_some() { + queue.push_back(u); + *queued_flag = true; + } + } + while let Some(u) = queue.pop_front() { + queued[u] = false; + let Some(du) = self.dist[u] else { + continue; + }; + for &(v, weight) in &self.adj[u] { + let cand = du + weight; + let better = self.dist[v].is_none_or(|old| cand < old); + if better { + self.dist[v] = Some(cand); + self.pred[v] = Some(u); + if !queued[v] { + queued[v] = true; + queue.push_back(v); + } + } + } + } + } + + fn find_pivots(&mut self, bound: Option, sources: &[usize]) -> (Vec, Vec) { + let mut witnessed = FxHashSet::default(); + witnessed.extend(sources.iter().copied()); + let mut layer: Vec = sources.to_vec(); + + for _ in 0..self.k { + let mut next = Vec::new(); + for &u in &layer { + for &(v, weight) in &self.adj[u] { + let Some(cand) = self.relax(u, v, weight) else { + continue; + }; + if less_than(&cand, bound) { + next.push(v); + witnessed.insert(v); + } + } + } + if witnessed.len() > self.k.saturating_mul(sources.len()) { + return (sources.to_vec(), witnessed.into_iter().collect()); + } + if next.is_empty() { + break; + } + layer = next; + } + + let in_w: FxHashSet = witnessed.iter().copied().collect(); + let mut children: FxHashMap> = FxHashMap::default(); + let mut has_parent = FxHashSet::default(); + for &v in &witnessed { + if let Some(u) = self.pred[v].filter(|u| in_w.contains(u)) { + children.entry(u).or_default().push(v); + has_parent.insert(v); + } + } + + let mut memo = FxHashMap::default(); + let mut pivots = Vec::new(); + for &u in sources { + if !in_w.contains(&u) || has_parent.contains(&u) { + continue; + } + if tree_size(u, &children, &mut memo) >= self.k { + pivots.push(u); + } + } + (pivots, witnessed.into_iter().collect()) + } +} + +fn tree_size( + u: usize, + children: &FxHashMap>, + memo: &mut FxHashMap, +) -> usize { + if let Some(&sz) = memo.get(&u) { + return sz; + } + let mut total = 1; + if let Some(vs) = children.get(&u) { + for &v in vs { + total += tree_size(v, children, memo); + } + } + memo.insert(u, total); + total +} + +/// Partial-order frontier: insert, batch-prepend, and pull the next `m` keys. +struct PartialQueue { + by_value: BTreeMap>, + values: FxHashMap, + bound: Option, + m: usize, +} + +impl PartialQueue { + fn new(bound: Option, m: usize) -> Self { + Self { + by_value: BTreeMap::new(), + values: FxHashMap::default(), + bound, + m, + } + } + + fn is_empty(&self) -> bool { + self.values.is_empty() + } + + fn insert(&mut self, key: usize, value: C) { + if !less_than(&value, self.bound) { + return; + } + if let Some(&old) = self.values.get(&key) { + if value >= old { + return; + } + self.remove_from_bucket(old, key); + } + self.values.insert(key, value); + self.by_value.entry(value).or_default().push(key); + } + + fn batch_prepend(&mut self, items: &[(usize, C)]) { + for &(key, value) in items { + self.insert(key, value); + } + } + + fn pull(&mut self) -> (Vec, Option) { + let mut taken = Vec::new(); + while taken.len() < self.m { + let Some((&value, _)) = self.by_value.first_key_value() else { + break; + }; + let mut bucket = self.by_value.remove(&value).unwrap_or_default(); + while let Some(key) = bucket.pop() { + if self.values.get(&key) != Some(&value) { + continue; + } + self.values.remove(&key); + taken.push(key); + if taken.len() == self.m { + if !bucket.is_empty() { + self.by_value.insert(value, bucket); + } + let sep = self + .by_value + .first_key_value() + .map(|(&v, _)| v) + .or(self.bound); + return (taken, sep); + } + } + } + let sep = self + .by_value + .first_key_value() + .map(|(&v, _)| v) + .or(self.bound); + (taken, sep) + } + + fn remove_from_bucket(&mut self, value: C, key: usize) { + if let Some(bucket) = self.by_value.get_mut(&value) { + if let Some(i) = bucket.iter().position(|&k| k == key) { + bucket.swap_remove(i); + } + if bucket.is_empty() { + self.by_value.remove(&value); + } + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 4f326e66..68bac404 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -27,6 +27,7 @@ //! - [paths counting](directed/count_paths/index.html): count the paths to the destination in an acyclic graph //! - [strongly connected components](directed/strongly_connected_components/index.html): find strongly connected components in a directed graph ([⇒ Wikipedia][Strongly connected components]) //! - [topological sorting](directed/topological_sort/index.html): find an acceptable topological order in a directed graph ([⇒ Wikipedia][Topological sorting]) +//! - [SSSP](directed/sssp/index.html): single-source shortest paths via BMSSP ([⇒ arXiv:2504.17033][BMSSP]) //! - [Yen](directed/yen/index.html): find k-shortest paths using Dijkstra ([⇒ Wikipedia][Yen]) //! //! ### Undirected graphs @@ -108,6 +109,7 @@ //! [Strongly connected components]: https://en.wikipedia.org/wiki/Strongly_connected_component //! [Topological sorting]: https://en.wikipedia.org/wiki/Topological_sorting //! [Yen]: https://en.wikipedia.org/wiki/Yen's_algorithm +//! [BMSSP]: https://arxiv.org/abs/2504.17033 use deprecate_until::deprecate_until; pub use num_traits; @@ -141,6 +143,7 @@ pub mod prelude { pub use crate::directed::fringe::*; pub use crate::directed::idastar::*; pub use crate::directed::iddfs::*; + pub use crate::directed::sssp::*; pub use crate::directed::strongly_connected_components::*; pub use crate::directed::topological_sort::*; pub use crate::directed::yen::*; diff --git a/tests/sssp.rs b/tests/sssp.rs new file mode 100644 index 00000000..b332fcea --- /dev/null +++ b/tests/sssp.rs @@ -0,0 +1,105 @@ +use pathfinding::prelude::{build_path, dijkstra, dijkstra_all, sssp, sssp_all}; +use rand::{rngs, RngExt as _}; +use std::collections::HashMap; + +#[expect(clippy::trivially_copy_pass_by_ref)] +fn successors(&n: &u32) -> Vec<(u32, usize)> { + if n <= 4 { + vec![(n * 2, 10), (n * 2 + 1, 10)] + } else { + vec![] + } +} + +#[test] +fn matches_dijkstra_all_on_small_tree() { + let sssp_map = sssp_all(&1, successors); + let dijkstra_map = dijkstra_all(&1, successors); + assert_eq!(sssp_map.len(), dijkstra_map.len()); + for (node, (_, cost)) in &dijkstra_map { + assert_eq!(sssp_map[node].1, *cost, "cost mismatch at {node}"); + } +} + +#[test] +fn sssp_path_matches_dijkstra() { + let (path, cost) = sssp(&1, successors, |&n| n == 9).unwrap(); + let (d_path, d_cost) = dijkstra(&1, successors, |&n| n == 9).unwrap(); + assert_eq!(cost, d_cost); + assert_eq!(path, d_path); + assert_eq!(build_path(&9, &sssp_all(&1, successors)), path); +} + +#[test] +fn start_is_goal() { + let (path, cost) = sssp(&1, successors, |&n| n == 1).unwrap(); + assert_eq!(path, vec![1]); + assert_eq!(cost, 0); +} + +#[test] +fn unreachable_goal() { + assert!(sssp(&1, successors, |&n| n == 100).is_none()); +} + +fn build_network(size: usize) -> pathfinding::prelude::Matrix { + let mut network = pathfinding::prelude::Matrix::new(size, size, 0); + let mut rng = rngs::ThreadRng::default(); + for a in 0..size { + for b in 0..size { + if rng.random_ratio(2, 3) { + network[(a, b)] = rng.random::() as usize + 1; + } + } + } + network +} + +fn neighbours( + network: pathfinding::prelude::Matrix, +) -> impl FnMut(&usize) -> Vec<(usize, usize)> { + move |&a| { + (0..network.rows) + .filter_map(|b| match network[(a, b)] { + 0 => None, + p => Some((b, p)), + }) + .collect() + } +} + +#[test] +fn random_graphs_match_dijkstra_costs() { + const SIZE: usize = 40; + let network = build_network(SIZE); + for start in 0..SIZE { + let sssp_map = sssp_all(&start, neighbours(network.clone())); + let dijkstra_map = dijkstra_all(&start, neighbours(network.clone())); + assert_eq!( + costs_only(&sssp_map), + costs_only(&dijkstra_map), + "costs differ from start {start} in {network:?}" + ); + } +} + +fn costs_only(map: &HashMap) -> HashMap { + map.iter().map(|(n, (_, c))| (n.clone(), *c)).collect() +} + +#[test] +fn grid_matches_dijkstra() { + let successors = |&(x, y): &(i32, i32)| { + [(1, 0), (-1, 0), (0, 1), (0, -1)] + .into_iter() + .map(move |(dx, dy)| ((x + dx, y + dy), 1_u32)) + .filter(|&((nx, ny), _)| (0..=8).contains(&nx) && (0..=8).contains(&ny)) + }; + let sssp_map = sssp_all(&(0, 0), successors); + let dijkstra_map = dijkstra_all(&(0, 0), successors); + assert_eq!(costs_only(&sssp_map), costs_only(&dijkstra_map)); + let (path, cost) = sssp(&(0, 0), successors, |n| *n == (8, 8)).unwrap(); + assert_eq!(cost, 16); + assert_eq!(path.first(), Some(&(0, 0))); + assert_eq!(path.last(), Some(&(8, 8))); +}