1use super::*;
2use std::{
3 cmp::{Ordering, Reverse},
4 collections::{BinaryHeap, VecDeque},
5 marker::PhantomData,
6 ops::{Add, Mul},
7};
8
9pub trait ShortestPathSemiRing {
10 type T: Clone + Ord;
11 fn source() -> Self::T;
12 fn inf() -> Self::T;
13 fn mul(x: &Self::T, y: &Self::T) -> Self::T;
14 fn add_assign(x: &mut Self::T, y: &Self::T) -> bool;
15}
16
17pub struct StandardSp<M>(PhantomData<fn() -> M>);
18impl<M> ShortestPathSemiRing for StandardSp<M>
19where
20 M: Monoid<T: Bounded + Ord>,
21{
22 type T = M::T;
23 fn source() -> Self::T {
24 M::unit()
25 }
26 fn inf() -> Self::T {
27 M::T::maximum()
28 }
29 fn mul(x: &Self::T, y: &Self::T) -> Self::T {
30 M::operate(x, y)
31 }
32 fn add_assign(x: &mut Self::T, y: &Self::T) -> bool {
33 if &*x > y {
34 *x = y.clone();
35 true
36 } else {
37 false
38 }
39 }
40}
41
42pub struct OptionSp<M>(PhantomData<fn() -> M>);
43impl<M> ShortestPathSemiRing for OptionSp<M>
44where
45 M: Monoid<T: Ord>,
46{
47 type T = Option<M::T>;
48 fn source() -> Self::T {
49 Some(M::unit())
50 }
51 fn inf() -> Self::T {
52 None
53 }
54 fn mul(x: &Self::T, y: &Self::T) -> Self::T {
55 match (x, y) {
56 (Some(x), Some(y)) => Some(M::operate(x, y)),
57 _ => None,
58 }
59 }
60 fn add_assign(x: &mut Self::T, y: &Self::T) -> bool {
61 if let Some(y) = y {
62 if let Some(x) = x {
63 if &*x > y {
64 *x = y.clone();
65 true
66 } else {
67 false
68 }
69 } else {
70 *x = Some(y.clone());
71 true
72 }
73 } else {
74 false
75 }
76 }
77}
78
79pub struct PathFoldingSp<M, S>(PhantomData<fn() -> (M, S)>);
80impl<M, S> ShortestPathSemiRing for PathFoldingSp<M, S>
81where
82 M: Monoid<T: Bounded + Ord>,
83 S: SemiRing,
84{
85 type T = PartialIgnoredOrd<M::T, S::T>;
86 fn source() -> Self::T {
87 PartialIgnoredOrd(M::unit(), S::one())
88 }
89 fn inf() -> Self::T {
90 PartialIgnoredOrd(M::T::maximum(), S::zero())
91 }
92 fn mul(x: &Self::T, y: &Self::T) -> Self::T {
93 PartialIgnoredOrd(M::operate(&x.0, &y.0), S::mul(&x.1, &y.1))
94 }
95 fn add_assign(x: &mut Self::T, y: &Self::T) -> bool {
96 match x.0.cmp(&y.0) {
97 Ordering::Equal => {
98 x.1 = S::add(&x.1, &y.1);
99 false
100 }
101 Ordering::Greater => {
102 *x = y.clone();
103 true
104 }
105 _ => false,
106 }
107 }
108}
109
110pub trait ParentPolicy<G>
111where
112 G: GraphBase,
113{
114 type State;
115 fn init(graph: &G) -> Self::State;
116 fn save_parent(graph: &G, state: &mut Self::State, from: G::VIndex, to: G::VIndex);
117}
118
119pub enum NoParent {}
120impl<G> ParentPolicy<G> for NoParent
121where
122 G: GraphBase,
123{
124 type State = ();
125 fn init(_graph: &G) {}
126 fn save_parent(_graph: &G, _state: &mut Self::State, _from: G::VIndex, _to: G::VIndex) {}
127}
128
129pub struct RecordParent;
130impl<G> ParentPolicy<G> for RecordParent
131where
132 G: GraphBase + VertexMap<Option<<G as GraphBase>::VIndex>>,
133{
134 type State = <G as VertexMap<Option<<G as GraphBase>::VIndex>>>::Vmap;
135 fn init(graph: &G) -> Self::State {
136 graph.construct_vmap(|| None)
137 }
138 fn save_parent(graph: &G, state: &mut Self::State, from: G::VIndex, to: G::VIndex) {
139 *graph.vmap_get_mut(state, to) = Some(from);
140 }
141}
142
143pub struct ShortestPathWithParent<G, S, P = RecordParent>
144where
145 G: GraphBase + VertexMap<S::T>,
146 S: ShortestPathSemiRing,
147 P: ParentPolicy<G>,
148{
149 pub dist: <G as VertexMap<S::T>>::Vmap,
150 pub parent: P::State,
151}
152
153impl<G, S> ShortestPathWithParent<G, S, RecordParent>
154where
155 G: GraphBase + VertexMap<S::T> + VertexMap<Option<<G as GraphBase>::VIndex>>,
156 S: ShortestPathSemiRing,
157{
158 pub fn path_to(&self, graph: &G, target: G::VIndex) -> Option<Vec<G::VIndex>> {
159 let dist: &S::T = graph.vmap_get(&self.dist, target);
160 if dist == &S::inf() {
161 return None;
162 }
163 let mut cur = target;
164 let mut path = vec![cur];
165 while let &Some(p) = graph.vmap_get(&self.parent, cur) {
166 path.push(p);
167 cur = p;
168 }
169 path.reverse();
170 Some(path)
171 }
172}
173
174pub trait ShortestPathExt: GraphBase {
175 fn standard_sp<'a, M>(&'a self) -> ShortestPathBuilder<'a, Self, StandardSp<M>>
176 where
177 Self: Sized + GraphBase,
178 M: Monoid<T: Bounded + Ord>,
179 {
180 ShortestPathBuilder {
181 graph: self,
182 _marker: PhantomData,
183 }
184 }
185
186 fn standard_sp_additive<'a, T>(
187 &'a self,
188 ) -> ShortestPathBuilder<'a, Self, StandardSp<AdditiveOperation<T>>>
189 where
190 Self: Sized + GraphBase,
191 T: Clone + Zero + Add<Output = T> + Bounded + Ord,
192 {
193 ShortestPathBuilder {
194 graph: self,
195 _marker: PhantomData,
196 }
197 }
198
199 fn option_sp<'a, M>(&'a self) -> ShortestPathBuilder<'a, Self, OptionSp<M>>
200 where
201 Self: Sized + GraphBase,
202 M: Monoid<T: Ord>,
203 {
204 ShortestPathBuilder {
205 graph: self,
206 _marker: PhantomData,
207 }
208 }
209
210 fn option_sp_additive<'a, T>(
211 &'a self,
212 ) -> ShortestPathBuilder<'a, Self, OptionSp<AdditiveOperation<T>>>
213 where
214 Self: Sized + GraphBase,
215 T: Clone + Zero + Add<Output = T> + Ord,
216 {
217 ShortestPathBuilder {
218 graph: self,
219 _marker: PhantomData,
220 }
221 }
222
223 fn path_folding_sp<'a, M, S>(&'a self) -> ShortestPathBuilder<'a, Self, PathFoldingSp<M, S>>
224 where
225 Self: Sized + GraphBase,
226 M: Monoid<T: Bounded + Ord>,
227 S: SemiRing,
228 {
229 ShortestPathBuilder {
230 graph: self,
231 _marker: PhantomData,
232 }
233 }
234
235 fn path_folding_sp_additive_addmul<'a, T, U>(
236 &'a self,
237 ) -> ShortestPathBuilder<'a, Self, PathFoldingSp<AdditiveOperation<T>, AddMulOperation<U>>>
238 where
239 Self: Sized + GraphBase,
240 T: Clone + Zero + Add<Output = T> + Bounded + Ord,
241 U: Clone + Zero + One + Add<Output = U> + Mul<Output = U>,
242 {
243 ShortestPathBuilder {
244 graph: self,
245 _marker: PhantomData,
246 }
247 }
248}
249impl<G> ShortestPathExt for G where G: GraphBase {}
250
251pub struct ShortestPathBuilder<'a, G, S, P = NoParent>
252where
253 G: GraphBase,
254 S: ShortestPathSemiRing,
255 P: ParentPolicy<G>,
256{
257 graph: &'a G,
258 _marker: PhantomData<fn() -> (S, P)>,
259}
260
261impl<'a, G, S, P> ShortestPathBuilder<'a, G, S, P>
262where
263 G: GraphBase,
264 S: ShortestPathSemiRing,
265 P: ParentPolicy<G>,
266{
267 fn bfs_distance_core<M, I>(&self, sources: I, weight: &'a M) -> ShortestPathWithParent<G, S, P>
268 where
269 G: VertexMap<S::T> + AdjacencyView<'a, M, S::T>,
270 I: IntoIterator<Item = G::VIndex>,
271 {
272 let graph = self.graph;
273 let mut dist = graph.construct_vmap(S::inf);
274 let mut parent = P::init(graph);
275 let mut deq = VecDeque::new();
276 for source in sources.into_iter() {
277 *graph.vmap_get_mut(&mut dist, source) = S::source();
278 deq.push_back(source);
279 }
280 let zero = S::source();
281 while let Some(u) = deq.pop_front() {
282 for a in graph.aviews(weight, u) {
283 let v = a.vindex();
284 let w = a.avalue();
285 let nd = S::mul(graph.vmap_get(&dist, u), &w);
286 if S::add_assign(graph.vmap_get_mut(&mut dist, v), &nd) {
287 P::save_parent(graph, &mut parent, u, v);
288 if w == zero {
289 deq.push_front(v);
290 } else {
291 deq.push_back(v);
292 }
293 }
294 }
295 }
296 ShortestPathWithParent { dist, parent }
297 }
298
299 fn dijkstra_core<M, I>(&self, sources: I, weight: &'a M) -> ShortestPathWithParent<G, S, P>
300 where
301 G: VertexMap<S::T> + AdjacencyView<'a, M, S::T>,
302 I: IntoIterator<Item = G::VIndex>,
303 {
304 let graph = self.graph;
305 let mut dist = graph.construct_vmap(S::inf);
306 let mut parent = P::init(graph);
307 let mut heap = BinaryHeap::new();
308 for source in sources.into_iter() {
309 *graph.vmap_get_mut(&mut dist, source) = S::source();
310 heap.push(PartialIgnoredOrd(Reverse(S::source()), source));
311 }
312 while let Some(PartialIgnoredOrd(Reverse(d), u)) = heap.pop() {
313 if graph.vmap_get(&dist, u) != &d {
314 continue;
315 }
316 let d = graph.vmap_get(&dist, u).clone();
317 for a in graph.aviews(weight, u) {
318 let v = a.vindex();
319 let nd = S::mul(&d, &a.avalue());
320 if S::add_assign(graph.vmap_get_mut(&mut dist, v), &nd) {
321 P::save_parent(graph, &mut parent, u, v);
322 heap.push(PartialIgnoredOrd(Reverse(nd), v));
323 }
324 }
325 }
326 ShortestPathWithParent { dist, parent }
327 }
328
329 fn bellman_ford_core<M, I>(
330 &self,
331 sources: I,
332 weight: &'a M,
333 check: bool,
334 ) -> Option<ShortestPathWithParent<G, S, P>>
335 where
336 G: Vertices + VertexMap<S::T> + AdjacencyView<'a, M, S::T> + VertexSize,
337 I: IntoIterator<Item = G::VIndex>,
338 P: ParentPolicy<G>,
339 {
340 let graph = self.graph;
341 let mut dist = graph.construct_vmap(S::inf);
342 let mut parent = P::init(graph);
343 for source in sources.into_iter() {
344 *graph.vmap_get_mut(&mut dist, source) = S::source();
345 }
346 let vsize = graph.vsize();
347 for _ in 1..vsize {
348 let mut updated = false;
349 for u in graph.vertices() {
350 for a in graph.aviews(weight, u) {
351 let v = a.vindex();
352 let nd = S::mul(graph.vmap_get(&dist, u), &a.avalue());
353 if S::add_assign(graph.vmap_get_mut(&mut dist, v), &nd) {
354 P::save_parent(graph, &mut parent, u, v);
355 updated = true;
356 }
357 }
358 }
359 if !updated {
360 break;
361 }
362 }
363 if check {
364 for u in graph.vertices() {
365 for a in graph.aviews(weight, u) {
366 let v = a.vindex();
367 let nd = S::mul(graph.vmap_get(&dist, u), &a.avalue());
368 if S::add_assign(graph.vmap_get_mut(&mut dist, v), &nd) {
369 return None;
370 }
371 }
372 }
373 }
374 Some(ShortestPathWithParent { dist, parent })
375 }
376}
377
378impl<'a, G, S> ShortestPathBuilder<'a, G, S, NoParent>
379where
380 G: GraphBase,
381 S: ShortestPathSemiRing,
382{
383 pub fn with_parent(self) -> ShortestPathBuilder<'a, G, S, RecordParent>
384 where
385 G: VertexMap<Option<<G as GraphBase>::VIndex>>,
386 {
387 ShortestPathBuilder {
388 graph: self.graph,
389 _marker: PhantomData,
390 }
391 }
392
393 pub fn bfs_distance<M, I>(&self, sources: I, weight: &'a M) -> <G as VertexMap<S::T>>::Vmap
394 where
395 G: VertexMap<S::T> + AdjacencyView<'a, M, S::T>,
396 I: IntoIterator<Item = G::VIndex>,
397 {
398 self.bfs_distance_core::<M, I>(sources, weight).dist
399 }
400
401 pub fn dijkstra<M, I>(&self, sources: I, weight: &'a M) -> <G as VertexMap<S::T>>::Vmap
402 where
403 G: VertexMap<S::T> + AdjacencyView<'a, M, S::T>,
404 I: IntoIterator<Item = G::VIndex>,
405 {
406 self.dijkstra_core::<M, I>(sources, weight).dist
407 }
408
409 pub fn bellman_ford<M, I>(
410 &self,
411 sources: I,
412 weight: &'a M,
413 check: bool,
414 ) -> Option<<G as VertexMap<S::T>>::Vmap>
415 where
416 G: Vertices + VertexMap<S::T> + AdjacencyView<'a, M, S::T> + VertexSize,
417 I: IntoIterator<Item = G::VIndex>,
418 {
419 self.bellman_ford_core::<M, I>(sources, weight, check)
420 .map(|sp| sp.dist)
421 }
422
423 pub fn warshall_floyd_ap<M>(
424 &self,
425 weight: &'a M,
426 ) -> <G as VertexMap<<G as VertexMap<S::T>>::Vmap>>::Vmap
427 where
428 G: Vertices
429 + VertexMap<S::T, Vmap: Clone>
430 + VertexMap<<G as VertexMap<S::T>>::Vmap>
431 + AdjacencyView<'a, M, S::T>,
432 {
433 let graph = self.graph;
434 let mut dist = graph.construct_vmap(|| graph.construct_vmap(S::inf));
435 for u in graph.vertices() {
436 *graph.vmap_get_mut(graph.vmap_get_mut(&mut dist, u), u) = S::source();
437 }
438 for u in graph.vertices() {
439 for a in graph.aviews(weight, u) {
440 S::add_assign(
441 graph.vmap_get_mut(graph.vmap_get_mut(&mut dist, u), a.vindex()),
442 &a.avalue(),
443 );
444 }
445 }
446 for k in graph.vertices() {
447 for i in graph.vertices() {
448 for j in graph.vertices() {
449 let d1 = graph.vmap_get(graph.vmap_get(&dist, i), k);
450 let d2 = graph.vmap_get(graph.vmap_get(&dist, k), j);
451 let nd = S::mul(d1, d2);
452 S::add_assign(graph.vmap_get_mut(graph.vmap_get_mut(&mut dist, i), j), &nd);
453 }
454 }
455 }
456 dist
457 }
458}
459
460impl<'a, G, S> ShortestPathBuilder<'a, G, S, RecordParent>
461where
462 G: GraphBase + VertexMap<Option<<G as GraphBase>::VIndex>>,
463 S: ShortestPathSemiRing,
464{
465 pub fn bfs_distance<M, I>(&self, sources: I, weight: &'a M) -> ShortestPathWithParent<G, S>
466 where
467 G: VertexMap<S::T> + AdjacencyView<'a, M, S::T>,
468 I: IntoIterator<Item = G::VIndex>,
469 {
470 self.bfs_distance_core::<M, I>(sources, weight)
471 }
472
473 pub fn dijkstra<M, I>(&self, sources: I, weight: &'a M) -> ShortestPathWithParent<G, S>
474 where
475 G: VertexMap<S::T> + AdjacencyView<'a, M, S::T>,
476 I: IntoIterator<Item = G::VIndex>,
477 {
478 self.dijkstra_core::<M, I>(sources, weight)
479 }
480
481 pub fn bellman_ford<M, I>(
482 &self,
483 sources: I,
484 weight: &'a M,
485 check: bool,
486 ) -> Option<ShortestPathWithParent<G, S>>
487 where
488 G: Vertices + VertexMap<S::T> + AdjacencyView<'a, M, S::T> + VertexSize,
489 I: IntoIterator<Item = G::VIndex>,
490 {
491 self.bellman_ford_core::<M, I>(sources, weight, check)
492 }
493}
494
495#[cfg(test)]
496mod tests {
497 use super::*;
498 use crate::{
499 num::{Saturating, mint_basic::MInt998244353},
500 rand,
501 tools::{PartialOrdExt, Xorshift},
502 };
503 use std::collections::HashMap;
504
505 #[test]
506 fn test_shortest_path() {
507 let mut rng = Xorshift::default();
508 for _ in 0..100 {
509 rand!(rng, n: 1..100, m: 1..200, edges: [(0..n, 0..n); m], w: [0..100_000i64; m]);
510 let g = DirectedSparseGraph::from_edges(n, edges);
511 let dijkstra: Vec<_> = (0..n)
512 .map(|src| g.option_sp_additive().dijkstra([src], &|eid| Some(w[eid])))
513 .collect();
514 let bellman_ford: Vec<_> = (0..n)
515 .map(|src| {
516 g.option_sp_additive()
517 .bellman_ford([src], &|eid| Some(w[eid]), false)
518 .unwrap()
519 })
520 .collect();
521 let warshall_floyd = g
522 .option_sp_additive()
523 .warshall_floyd_ap(&|eid| Some(w[eid]));
524 assert_eq!(dijkstra, bellman_ford);
525 assert_eq!(dijkstra, warshall_floyd);
526 }
527 }
528
529 #[test]
530 fn test_spfa() {
531 let mut rng = Xorshift::default();
532 for _ in 0..100 {
533 rand!(rng, n: 1..100, m: 1..200, edges: [(0..n, 0..n); m], ub: 0..=1i64, w: [0..=ub * 100_000i64; m]);
534 let g = DirectedSparseGraph::from_edges(n, edges);
535 let bfs: Vec<_> = (0..n)
536 .map(|src| {
537 g.option_sp_additive()
538 .bfs_distance([src], &|eid| Some(w[eid]))
539 })
540 .collect();
541 let warshall_floyd = g
542 .option_sp_additive()
543 .warshall_floyd_ap(&|eid| Some(w[eid]));
544 assert_eq!(bfs, warshall_floyd);
545 }
546 }
547
548 #[test]
549 fn test_shortest_path_with_parent() {
550 let mut rng = Xorshift::default();
551 for _ in 0..100 {
552 rand!(rng, n: 1..100, m: 1..200, edges: [(0..n, 0..n); m], ub: 0..=1i64, w: [0..=ub * 100_000i64; m]);
553 let mut cost: HashMap<_, _> = HashMap::new();
554 for (&(u, v), &w) in edges.iter().zip(w.iter()) {
555 cost.entry((u, v)).or_insert(w).chmin(w);
556 }
557 let g = DirectedSparseGraph::from_edges(n, edges);
558 for src in 0..n {
559 let bfs = g
560 .option_sp_additive()
561 .with_parent()
562 .bfs_distance([src], &|eid| Some(w[eid]));
563 let dijkstra = g
564 .option_sp_additive()
565 .with_parent()
566 .dijkstra([src], &|eid| Some(w[eid]));
567 let bellman_ford = g
568 .option_sp_additive()
569 .with_parent()
570 .bellman_ford([src], &|eid| Some(w[eid]), false)
571 .unwrap();
572 let dist = g.option_sp_additive().dijkstra([src], &|eid| Some(w[eid]));
573 assert_eq!(bfs.dist, dist);
574 assert_eq!(dijkstra.dist, dist);
575 assert_eq!(bellman_ford.dist, dist);
576 for (target, &dist) in dist.iter().enumerate() {
577 match dist {
578 None => {
579 assert!(bfs.path_to(&g, target).is_none());
580 assert!(dijkstra.path_to(&g, target).is_none());
581 assert!(bellman_ford.path_to(&g, target).is_none());
582 }
583 Some(dist) => {
584 let path_bfs = bfs.path_to(&g, target).unwrap();
585 assert_eq!(*path_bfs.first().unwrap(), src);
586 assert_eq!(*path_bfs.last().unwrap(), target);
587 assert_eq!(
588 path_bfs
589 .windows(2)
590 .map(|w| cost[&(w[0], w[1])])
591 .sum::<i64>(),
592 dist
593 );
594 let path_dijkstra = dijkstra.path_to(&g, target).unwrap();
595 assert_eq!(*path_dijkstra.first().unwrap(), src);
596 assert_eq!(*path_dijkstra.last().unwrap(), target);
597 assert_eq!(
598 path_dijkstra
599 .windows(2)
600 .map(|w| cost[&(w[0], w[1])])
601 .sum::<i64>(),
602 dist
603 );
604 let path_bellman_ford = bellman_ford.path_to(&g, target).unwrap();
605 assert_eq!(*path_bellman_ford.first().unwrap(), src);
606 assert_eq!(*path_bellman_ford.last().unwrap(), target);
607 assert_eq!(
608 path_bellman_ford
609 .windows(2)
610 .map(|w| cost[&(w[0], w[1])])
611 .sum::<i64>(),
612 dist
613 );
614 }
615 }
616 }
617 }
618 }
619 }
620
621 #[test]
622 fn test_path_folding() {
623 let mut rng = Xorshift::default();
624 for _ in 0..100 {
625 rand!(rng, n: 1..100, m: 1..200, edges: [(0..n, 0..n); m], w: [0..100_000usize; m]);
626 let g = DirectedSparseGraph::from_edges(n, edges);
627 let dijkstra: Vec<_> = (0..n)
628 .map(|src| {
629 g.path_folding_sp_additive_addmul().dijkstra([src], &|eid| {
630 PartialIgnoredOrd(Saturating(w[eid]), MInt998244353::one())
631 })
632 })
633 .collect();
634 let bellman_ford: Vec<_> = (0..n)
635 .map(|src| {
636 g.path_folding_sp_additive_addmul()
637 .bellman_ford(
638 [src],
639 &|eid| PartialIgnoredOrd(Saturating(w[eid]), MInt998244353::one()),
640 false,
641 )
642 .unwrap()
643 })
644 .collect();
645 let warshall_floyd = g
646 .path_folding_sp_additive_addmul()
647 .warshall_floyd_ap(&|eid| {
648 PartialIgnoredOrd(Saturating(w[eid]), MInt998244353::one())
649 });
650 assert_eq!(dijkstra, bellman_ford);
651 assert_eq!(dijkstra, warshall_floyd);
652 }
653 }
654}