haar_lib/graph/
yen.rs

1//! 最短パスを`k`個列挙する。
2//!
3//! # Problems
4//! - <https://yukicoder.me/problems/no/1069>
5use std::ops::{Add, AddAssign};
6use std::{cmp::Reverse, collections::BinaryHeap};
7
8use crate::{graph::*, num::one_zero::Zero};
9
10type Path = Vec<usize>;
11
12fn shortest_path<D: Direction, E: EdgeTrait>(
13    g: &Graph<D, E>,
14    from: usize,
15    t: usize,
16    usable: &[bool],
17    valid: &[Vec<bool>],
18) -> Option<(E::Weight, Path)>
19where
20    E::Weight: Zero + Add<Output = E::Weight> + Ord + Eq + Copy,
21{
22    let n = g.len();
23    let mut visited = vec![false; n];
24    let mut dist = vec![None; n];
25    let mut restore = vec![(0, 0); n];
26    let mut pq = BinaryHeap::new();
27
28    dist[from] = Some(E::Weight::zero());
29    pq.push(Reverse((E::Weight::zero(), from)));
30
31    while let Some(Reverse((d, i))) = pq.pop() {
32        if visited[i] {
33            continue;
34        }
35        visited[i] = true;
36
37        for (k, e) in g.nodes[i].edges.iter().enumerate() {
38            if !valid[i][k] || !usable[e.to()] {
39                continue;
40            }
41
42            if dist[e.to()].is_none_or(|x| x > d + e.weight()) {
43                dist[e.to()] = Some(d + e.weight());
44                restore[e.to()] = (i, k);
45                if !visited[e.to()] {
46                    pq.push(Reverse((dist[e.to()].unwrap(), e.to())));
47                }
48            }
49        }
50    }
51
52    if let Some(d) = dist[t] {
53        let mut p = vec![];
54
55        let mut cur = t;
56        while cur != from {
57            let (i, j) = restore[cur];
58            p.push(j);
59            cur = i;
60        }
61
62        p.reverse();
63
64        Some((d, p))
65    } else {
66        None
67    }
68}
69
70/// 有向グラフ`g`上で`from`から`to`へのパスを、その距離が小さい順に`k`個を返す。
71pub fn yen_algorithm<D: Direction, E: EdgeTrait>(
72    g: &Graph<D, E>,
73    from: usize,
74    to: usize,
75    k: usize,
76) -> Vec<Option<(E::Weight, Path)>>
77where
78    E::Weight: Zero + Add<Output = E::Weight> + AddAssign + Ord + Eq + Copy,
79{
80    let n = g.len();
81    let mut result: Vec<Option<(E::Weight, Path)>> = vec![None; k];
82    let mut stock = BinaryHeap::new();
83    let mut valid = (0..n)
84        .map(|i| vec![true; g.nodes[i].edges.len()])
85        .collect::<Vec<_>>();
86
87    for i in 0..k {
88        if i == 0 {
89            let usable = vec![true; n];
90            if let Some((c, p)) = shortest_path(g, from, to, &usable, &valid) {
91                stock.push(Reverse((c, p)));
92            }
93        } else {
94            let mut prev_path = vec![];
95
96            let mut cur = from;
97            for &u in &result[i - 1].as_ref().unwrap().1 {
98                prev_path.push(cur);
99                cur = g.nodes[cur].edges[u].to();
100            }
101            prev_path.push(to);
102
103            let mut check = vec![true; i];
104            let mut usable = vec![true; n];
105
106            for k in 0..prev_path.len() - 1 {
107                let u = prev_path[k];
108
109                for j in 0..i {
110                    if check[j] {
111                        valid[u][result[j].as_ref().unwrap().1[k]] = false;
112                    }
113                }
114
115                if let Some((mut c, p)) = shortest_path(g, u, to, &usable, &valid) {
116                    let mut temp = vec![];
117
118                    for (j, &p) in prev_path.iter().enumerate().take(k) {
119                        let v = result[i - 1].as_ref().unwrap().1[j];
120                        c += g.nodes[p].edges[v].weight();
121                        temp.push(v);
122                    }
123
124                    temp.extend(p.into_iter());
125                    stock.push(Reverse((c, temp)));
126                }
127
128                usable[u] = false;
129
130                for j in 0..i {
131                    if check[j] {
132                        valid[u][result[j].as_ref().unwrap().1[k]] = true;
133                    }
134                }
135
136                for j in 0..i {
137                    if check[j]
138                        && prev_path[k + 1]
139                            != g.nodes[u].edges[result[j].as_ref().unwrap().1[k]].to()
140                    {
141                        check[j] = false;
142                    }
143                }
144            }
145        }
146
147        if stock.is_empty() {
148            break;
149        }
150
151        result[i] = stock.pop().map(|a| a.0);
152
153        while stock.peek().map(|a| &a.0) == result[i].as_ref() {
154            stock.pop();
155        }
156    }
157
158    result
159}
haar_lib/graph/yen.rs

haar_lib/graph/
yen.rs
