Skip to main content

nekolib/graph/
dijkstra_.rs

1//! 最短距離 (Dijkstra)。
2
3use std::cmp::Reverse;
4use std::collections::BinaryHeap;
5
6/// Dijkstra 法に基づく最短距離。
7///
8/// # Parameters
9/// - `n`: 頂点数
10/// - `s`: 始点
11/// - `zero`: 距離を表す型の $0$
12/// - `index`: 頂点から添字への番号づけをする関数
13/// - `delta`: 頂点 `v` と関数 `search` を受け取る関数
14///
15/// `delta` は、`v` の各隣接頂点 `nv` とそこへの距離 `ew` に対して、
16/// `search(nv, ew)` を呼び出す必要がある。
17///
18/// # Examples
19/// `g[v]` が `v` の隣接頂点を持つ、よくある隣接リストにおける例を載せる。
20/// 次のようなグラフである。
21///
22/// ```text
23///      2        3
24/// (0) ---> (1) ---> (2)
25///  ^        |        |
26///  | 1      |        | 4
27///  |        |  9     v
28/// (3)       +-----> (4)
29/// ```
30///
31/// ```
32/// use nekolib::graph::dijkstra;
33///
34/// let g = vec![
35///     vec![(1, 2)],
36///     vec![(2, 3), (4, 9)],
37///     vec![(4, 4)],
38///     vec![(0, 1)],
39///     vec![],
40/// ];
41/// let index = |&v: &usize| v;
42/// let delta = |&v: &usize| g[v].iter().cloned();
43/// let dist = dijkstra(5, 0, 0_i32, index, delta);
44///
45/// assert_eq!(dist, vec![Some(0), Some(2), Some(5), None, Some(9)]);
46/// ```
47///
48/// 頂点数の上限 $n$ さえ既知であれば、`index` は動的に決められる。
49/// ```
50/// use std::collections::BTreeMap;
51///
52/// use nekolib::graph::dijkstra;
53///
54/// let n = 10;
55///
56/// let s = 10001;
57/// let t = 10015;
58/// let u = t - 1;
59/// let mut enc = BTreeMap::new();
60/// let index = |&v: &usize| match enc.get(&v) {
61///     Some(&i) => i,
62///     None => {
63///         let len = enc.len();
64///         enc.insert(v, len);
65///         len
66///     }
67/// };
68/// let delta =
69///   |&v: &usize| Some((v + 2, 1)).filter(|&(v, _)| v <= t).into_iter();
70///
71/// let res = dijkstra(n, s, 0, index, delta);
72/// assert_eq!(res[enc[&t]], Some(7));
73/// assert!(!enc.contains_key(&u))
74/// ```
75///
76/// # Notes
77/// 複数のグラフを扱う際に `delta` を使い回すと、意図しないグラフを見てしまいがちなので注意。
78///
79/// # Complexity
80/// 二分ヒープを用いる実装なので、$O(|E|\\log(|V|))$ 時間。
81/// ここで、$V$ は頂点集合、$E$ は辺集合である。
82///
83/// # References
84/// - <https://niuez.github.io/posts/impl_abstract_dijkstra/>
85///     - これを意識していますが、辺の取得にイテレータを使っているので少々異なります。
86pub fn dijkstra<V, W, E>(
87    n: usize,
88    s: V,
89    zero: W,
90    mut index: impl FnMut(&V) -> usize,
91    mut delta: impl FnMut(&V) -> E,
92) -> Vec<Option<W>>
93where
94    V: Ord,
95    W: Ord + std::ops::Add<W, Output = W> + Clone,
96    E: Iterator<Item = (V, W)>,
97{
98    let mut dist: Vec<Option<W>> = vec![None; n];
99    let si = index(&s);
100    dist[si] = Some(zero);
101    let mut pq: BinaryHeap<_> =
102        vec![Reverse((dist[si].clone().unwrap(), s))].into();
103
104    while let Some(Reverse((w, v))) = pq.pop() {
105        match &dist[index(&v)] {
106            Some(cw) if cw < &w => continue,
107            _ => {}
108        }
109        for (nv, ew) in delta(&v) {
110            let nw = w.clone() + ew;
111            match &dist[index(&nv)] {
112                Some(cw) if cw <= &nw => {}
113                _ => {
114                    dist[index(&nv)] = Some(nw.clone());
115                    pq.push(Reverse((nw, nv)));
116                }
117            }
118        }
119    }
120    dist
121}