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}