use {
crate::{
Order,
OutNeighborsWeighted,
},
core::cmp::Reverse,
std::collections::BinaryHeap,
};
type Step = (usize, usize);
#[derive(Clone, Debug)]
pub struct DijkstraDist<'a, D> {
digraph: &'a D,
dist: Vec<usize>,
heap: BinaryHeap<(Reverse<usize>, usize)>,
}
impl<'a, D> DijkstraDist<'a, D>
where
D: Order,
{
#[must_use]
pub fn new<T>(digraph: &'a D, sources: T) -> Self
where
D: Order,
T: Iterator<Item = usize> + Clone,
{
let order = digraph.order();
let mut heap = BinaryHeap::with_capacity(order);
let mut dist = vec![usize::MAX; order];
let dist_ptr = dist.as_mut_ptr();
for u in sources {
unsafe { *dist_ptr.add(u) = 0 };
heap.push((Reverse(0), u));
}
Self {
digraph,
dist,
heap,
}
}
#[must_use]
pub fn distances(&mut self) -> Vec<usize>
where
D: OutNeighborsWeighted<Weight = usize>,
{
let mut dist = vec![usize::MAX; self.digraph.order()];
let ptr = dist.as_mut_ptr();
for u in self {
unsafe { *ptr.add(u.0) = u.1 };
}
dist
}
}
impl<D> Iterator for DijkstraDist<'_, D>
where
D: Order + OutNeighborsWeighted<Weight = usize>,
{
type Item = Step;
fn next(&mut self) -> Option<Self::Item> {
let (Reverse(w_prev), u) = self.heap.pop()?;
let dist_ptr = self.dist.as_mut_ptr();
for (v, w) in self.digraph.out_neighbors_weighted(u) {
let w_next = w_prev + w;
let dist_v = unsafe { *dist_ptr.add(v) };
if w_next < dist_v {
unsafe { *dist_ptr.add(v) = w_next };
self.heap.push((Reverse(w_next), v));
}
}
if w_prev == unsafe { *dist_ptr.add(u) } {
return Some((u, w_prev));
}
None
}
}
#[cfg(test)]
mod tests {
use {
super::*,
crate::repr::adjacency_list_weighted::fixture::{
bang_jensen_94_usize,
bang_jensen_96_usize,
kattis_bryr_1_usize,
kattis_bryr_2_usize,
kattis_bryr_3_usize,
kattis_crosscountry_usize,
kattis_shortestpath1_usize,
},
std::iter::once,
};
#[test]
fn iter_bang_jensen_94() {
let digraph = bang_jensen_94_usize();
assert!(DijkstraDist::new(&digraph, once(0)).eq([
(0, 0),
(2, 1),
(1, 1),
(5, 2),
(4, 2),
(3, 2),
(6, 3)
]));
}
#[test]
fn iter_bang_jensen_96() {
let digraph = bang_jensen_96_usize();
assert!(DijkstraDist::new(&digraph, once(0)).eq([
(0, 0),
(2, 3),
(4, 4),
(1, 5),
(3, 6),
(5, 7),
]));
}
#[test]
fn iter_kattis_bryr_1() {
let digraph = kattis_bryr_1_usize();
assert!(DijkstraDist::new(&digraph, once(0)).eq([
(0, 0),
(2, 1),
(1, 1),
]));
}
#[test]
fn iter_kattis_bryr_2() {
let digraph = kattis_bryr_2_usize();
assert!(DijkstraDist::new(&digraph, once(0)).eq([
(0, 0),
(3, 1),
(1, 1),
(4, 2),
(2, 2),
(5, 3),
]));
}
#[test]
fn iter_kattis_bryr_3() {
let digraph = kattis_bryr_3_usize();
assert!(DijkstraDist::new(&digraph, once(0)).eq([
(0, 0),
(3, 0),
(7, 0),
(5, 0),
(8, 0),
(4, 0),
(1, 0),
(9, 1),
(6, 1),
(2, 1)
]));
}
#[test]
fn iter_kattis_crosscountry() {
let digraph = kattis_crosscountry_usize();
assert!(DijkstraDist::new(&digraph, once(0)).eq([
(0, 0),
(1, 1),
(2, 3),
(3, 10)
]));
}
#[test]
fn iter_kattis_shortestpath1() {
let digraph = kattis_shortestpath1_usize();
assert!(DijkstraDist::new(&digraph, once(0)).eq([
(0, 0),
(1, 2),
(2, 4)
]));
}
#[test]
fn distances_bang_jensen_94() {
let digraph = bang_jensen_94_usize();
assert!(
DijkstraDist::new(&digraph, once(0))
.distances()
.iter()
.eq(&[0, 1, 1, 2, 2, 2, 3])
);
}
#[test]
fn distances_bang_jensen_96() {
let digraph = bang_jensen_96_usize();
assert!(
DijkstraDist::new(&digraph, once(0))
.distances()
.iter()
.eq(&[0, 5, 3, 6, 4, 7])
);
}
#[test]
fn distances_kattis_bryr_1() {
let digraph = kattis_bryr_1_usize();
assert!(
DijkstraDist::new(&digraph, once(0))
.distances()
.iter()
.eq(&[0, 1, 1])
);
}
#[test]
fn distances_kattis_bryr_2() {
let digraph = kattis_bryr_2_usize();
assert!(
DijkstraDist::new(&digraph, once(0))
.distances()
.iter()
.eq(&[0, 1, 2, 1, 2, 3])
);
}
#[test]
fn distances_kattis_bryr_3() {
let digraph = kattis_bryr_3_usize();
assert!(
DijkstraDist::new(&digraph, once(0))
.distances()
.iter()
.eq(&[0, 0, 1, 0, 0, 0, 1, 0, 0, 1])
);
}
#[test]
fn distances_kattis_crosscountry() {
let digraph = kattis_crosscountry_usize();
assert!(
DijkstraDist::new(&digraph, once(0))
.distances()
.iter()
.eq(&[0, 1, 3, 10])
);
}
#[test]
fn distances_kattis_shortestpath1() {
let digraph = kattis_shortestpath1_usize();
assert!(
DijkstraDist::new(&digraph, once(0))
.distances()
.iter()
.eq(&[0, 2, 4, usize::MAX])
);
}
}