competitive_library/graph/
dijkstra_restore_path.rs

1//! Dijkstra's algorithm with path restoration.
2use std::{cmp::Ordering, collections::BinaryHeap};
3
4#[derive(Debug, Clone, Eq)]
5pub struct Node {
6    position: usize,
7    cost: i64,
8    from: Option<usize>,
9}
10impl Node {
11    #[inline]
12    pub fn new(position: usize, cost: i64, from: Option<usize>) -> Self {
13        Node {
14            position,
15            cost,
16            from,
17        }
18    }
19}
20impl PartialEq for Node {
21    fn eq(&self, other: &Node) -> bool {
22        self.cost.eq(&other.cost)
23    }
24}
25impl PartialOrd for Node {
26    fn partial_cmp(&self, other: &Node) -> Option<Ordering> {
27        Some(self.cmp(other))
28    }
29}
30impl Ord for Node {
31    fn cmp(&self, other: &Self) -> Ordering {
32        other.cost.cmp(&self.cost)
33    }
34}
35
36pub fn dijkstra(
37    edge: &[Vec<(usize, i64)>],
38    start: usize,
39    end: usize,
40    vertex: usize,
41) -> Option<(i64, Vec<usize>)> {
42    let mut costs = vec![None; edge.len()];
43    let mut nodes = BinaryHeap::new();
44    let mut previous = vec![None; vertex];
45    nodes.push(Node::new(start, 0, None));
46
47    while let Some(Node {
48        position,
49        cost,
50        from,
51    }) = nodes.pop()
52    {
53        if costs[position].is_some() {
54            continue;
55        }
56
57        previous[position] = from;
58        costs[position] = Some(cost);
59        if position == end {
60            return Some((cost, restore_path(end, &previous)));
61        }
62
63        edge[position]
64            .iter()
65            .filter(|(to, c)| costs[*to].filter(|&d| d <= cost + c).is_none())
66            .for_each(|&(to, c)| {
67                nodes.push(Node::new(to, cost + c, Some(position)));
68            });
69    }
70    None
71}
72
73fn restore_path(end: usize, previous: &[Option<usize>]) -> Vec<usize> {
74    let mut buff = end;
75    let mut v = vec![buff];
76
77    while let Some(i) = previous[buff] {
78        buff = i;
79        v.push(buff);
80    }
81    v.reverse();
82    v
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    #[test]
89    fn test_dijkstra() {
90        let graph = vec![
91            vec![(2, 10), (1, 1)],
92            vec![(3, 2)],
93            vec![(1, 1), (3, 3), (4, 1)],
94            vec![(0, 7), (4, 2)],
95            vec![],
96        ];
97        let l = graph.len();
98        for (start, end, ans) in &[
99            (0, 1, Some((1, vec![0, 1]))),
100            (0, 3, Some((3, vec![0, 1, 3]))),
101            (3, 0, Some((7, vec![3, 0]))),
102            (0, 4, Some((5, vec![0, 1, 3, 4]))),
103            (4, 0, None),
104        ] {
105            match dijkstra(&graph, *start, *end, l) {
106                Some((a, b)) => {
107                    assert_eq!(a, ans.as_ref().unwrap().0);
108                    assert_eq!(b, ans.as_ref().unwrap().1);
109                }
110                None => assert!(ans.is_none()),
111            }
112        }
113    }
114
115    #[test]
116    fn chooses_the_shorter_indirect_path() {
117        let graph = vec![vec![(1, 10), (2, 1)], vec![], vec![(1, 1)]];
118        assert_eq!(
119            dijkstra(&graph, 0, 1, graph.len()),
120            Some((2, vec![0, 2, 1]))
121        );
122    }
123}
competitive_library/graph/dijkstra_restore_path.rs

competitive_library/graph/
dijkstra_restore_path.rs
