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}