haar_lib/graph/
dijkstra.rs

1//! 非負重み付き最短経路 (Dijkstra)
2
3use crate::graph::*;
4use crate::num::{one_zero::Zero, traits::Unsigned};
5use std::{cmp::Reverse, collections::BinaryHeap, ops::Add};
6
7/// ダイクストラ法
8pub struct Dijkstra<'a, W, E> {
9    dist: Vec<Option<W>>,
10    prev: Vec<Option<&'a E>>,
11}
12
13impl<'a, E> Dijkstra<'a, E::Weight, E>
14where
15    E: EdgeTrait,
16    E::Weight: Add<Output = E::Weight> + Copy + Ord + Zero + Unsigned,
17{
18    /// グラフ`g`上で、始点から各頂点への最短パスを求める。
19    ///
20    /// **Time complexity** $O((E + V) \log V)$
21    pub fn new(g: &'a Graph<impl Direction, E>, src: &[usize]) -> Self {
22        let zero = E::Weight::zero();
23        let n = g.len();
24        let mut dist = vec![None; n];
25        let mut heap = BinaryHeap::new();
26        let mut check = vec![false; n];
27        let mut prev = vec![None; n];
28
29        for &u in src {
30            dist[u] = Some(zero);
31            heap.push(Reverse((zero, u)));
32        }
33
34        while let Some(Reverse((d, u))) = heap.pop() {
35            if check[u] {
36                continue;
37            }
38            check[u] = true;
39
40            for e in g.nodes[u].edges.iter() {
41                let (to, cost) = (e.to(), e.weight());
42
43                match dist[to] {
44                    Some(d2) if d2 <= d + cost => {}
45                    _ => {
46                        let d = d + cost;
47                        dist[to] = Some(d);
48                        prev[to] = Some(e);
49                        if !check[to] {
50                            heap.push(Reverse((d, to)));
51                        }
52                    }
53                }
54            }
55        }
56
57        Self { dist, prev }
58    }
59
60    /// 最短距離の配列への参照を返す。
61    pub fn min_dist_table(&self) -> &[Option<E::Weight>] {
62        &self.dist
63    }
64
65    /// `to`への最短距離を返す。
66    /// 到達不可能ならば、`None`を返す。
67    pub fn min_dist_to(&self, to: usize) -> Option<E::Weight> {
68        self.dist[to]
69    }
70
71    /// `to`への最短パスを返す。
72    /// 到達不可能ならば、`None`を返す。
73    pub fn min_path_to(&self, mut to: usize) -> Option<Vec<&E>> {
74        self.dist[to]?;
75        let mut ret = vec![];
76        loop {
77            let Some(e) = self.prev[to] else { break };
78            ret.push(e);
79            to = e.from();
80        }
81        ret.reverse();
82        Some(ret)
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn test_grl_1_a() {
92        // https://onlinejudge.u-aizu.ac.jp/courses/library/5/GRL/1/GRL_1_A
93
94        // sample 1
95        let mut graph = Graph::<Directed, Edge<u32, ()>>::new(4);
96        graph.extend(
97            vec![(0, 1, 1), (0, 2, 4), (1, 2, 2), (2, 3, 1), (1, 3, 5)]
98                .into_iter()
99                .map(|(u, v, w)| Edge::new(u, v, w, ())),
100        );
101        let ans = Dijkstra::new(&graph, &[0]);
102
103        assert_eq!(ans.min_dist_table(), [Some(0), Some(1), Some(3), Some(4)]);
104
105        // sample 2
106        let mut graph = Graph::<Directed, Edge<u32, ()>>::new(4);
107        graph.extend(
108            vec![
109                (0, 1, 1),
110                (0, 2, 4),
111                (2, 0, 1),
112                (1, 2, 2),
113                (3, 1, 1),
114                (3, 2, 5),
115            ]
116            .into_iter()
117            .map(|(u, v, w)| Edge::new(u, v, w, ())),
118        );
119        let ans = Dijkstra::new(&graph, &[1]);
120
121        assert_eq!(ans.min_dist_table(), [Some(3), Some(0), Some(2), None]);
122    }
123}