dijkstra/
lib.rs

1use std::{collections::BinaryHeap, ops::Add};
2
3pub struct Cert<V>(Vec<Option<V>>);
4pub struct NoCert;
5
6pub struct DijkstraSssp<V, W, I, C> {
7    cost: Vec<Option<W>>,
8    prev: C,
9    index: I,
10    src: V,
11}
12
13#[derive(Eq, PartialEq)]
14struct RevFst<F, S>(F, S);
15
16impl<F: Ord, S: Eq> Ord for RevFst<F, S> {
17    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
18        self.0.cmp(&other.0).reverse()
19    }
20}
21
22impl<F: Ord, S: Eq> PartialOrd for RevFst<F, S> {
23    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
24        Some(self.cmp(other))
25    }
26}
27
28impl<V, W, I> DijkstraSssp<V, W, I, Cert<V>>
29where
30    V: Eq + Clone,
31    W: Add<Output = W> + Ord + Clone,
32    I: Fn(&V) -> usize,
33{
34    pub fn new_cert<D, J>(
35        src: V,
36        len: usize,
37        zero: W,
38        index: I,
39        delta: D,
40    ) -> Self
41    where
42        D: Fn(&V) -> J,
43        J: Iterator<Item = (V, W)>,
44    {
45        let mut cost = vec![None; len];
46        let mut prev = vec![None; len];
47        let mut heap = BinaryHeap::new();
48        cost[index(&src)] = Some(zero.clone());
49        heap.push(RevFst(zero, src.clone()));
50        while let Some(RevFst(w, v)) = heap.pop() {
51            if let Some(cur_w) = &cost[index(&v)] {
52                if cur_w < &w {
53                    continue;
54                }
55            }
56            for (nv, dw) in delta(&v) {
57                let nw = w.clone() + dw;
58                let ni = index(&nv);
59                match &cost[ni] {
60                    Some(cur_w) if cur_w <= &nw => {}
61                    _ => {
62                        cost[ni] = Some(nw.clone());
63                        prev[ni] = Some(v.clone());
64                        heap.push(RevFst(nw, nv));
65                    }
66                }
67            }
68        }
69
70        Self { src, cost, prev: Cert(prev), index }
71    }
72    pub fn path(&self, dst: &V) -> Option<std::vec::IntoIter<V>> {
73        let mut i = (self.index)(dst);
74        if self.prev.0[i].is_none() {
75            return (&self.src == dst).then(|| vec![dst.clone()].into_iter());
76        }
77
78        let mut res = vec![dst.clone()];
79        while let Some(v) = &self.prev.0[i] {
80            i = (self.index)(v);
81            res.push(v.clone());
82        }
83        res.reverse();
84        Some(res.into_iter())
85    }
86}
87
88impl<V, W, I> DijkstraSssp<V, W, I, NoCert>
89where
90    V: Eq + Clone,
91    W: Add<Output = W> + Ord + Clone,
92    I: Fn(&V) -> usize,
93{
94    pub fn new<D, J>(src: V, len: usize, zero: W, index: I, delta: D) -> Self
95    where
96        D: Fn(&V) -> J,
97        J: Iterator<Item = (V, W)>,
98    {
99        let mut cost = vec![None; len];
100        let mut heap = BinaryHeap::new();
101        cost[index(&src)] = Some(zero.clone());
102        heap.push(RevFst(zero, src.clone()));
103        while let Some(RevFst(w, v)) = heap.pop() {
104            if let Some(cur_w) = &cost[index(&v)] {
105                if cur_w < &w {
106                    continue;
107                }
108            }
109            for (nv, dw) in delta(&v) {
110                let nw = w.clone() + dw;
111                let ni = index(&nv);
112                match &cost[ni] {
113                    Some(cur_w) if cur_w <= &nw => {}
114                    _ => {
115                        cost[ni] = Some(nw.clone());
116                        heap.push(RevFst(nw, nv));
117                    }
118                }
119            }
120        }
121
122        Self { src, cost, prev: NoCert, index }
123    }
124}
125
126impl<V, W, I, C> DijkstraSssp<V, W, I, C>
127where
128    V: Eq + Clone,
129    W: Add<Output = W> + Ord + Clone,
130    I: Fn(&V) -> usize,
131{
132    pub fn cost(&self, dst: &V) -> Option<W> {
133        self.cost[(self.index)(dst)].clone()
134    }
135}