competitive_library/graph/
dijkstra.rs

1//! ダイクストラ
2
3use std::{cmp::Ordering, collections::BinaryHeap};
4
5#[derive(Debug, Clone, Eq)]
6pub struct Node {
7    position: usize,
8    cost: i64,
9}
10impl Node {
11    #[inline]
12    pub fn new(position: usize, cost: i64) -> Self {
13        Node { position, cost }
14    }
15}
16impl PartialEq for Node {
17    fn eq(&self, other: &Node) -> bool {
18        self.cost.eq(&other.cost)
19    }
20}
21impl PartialOrd for Node {
22    fn partial_cmp(&self, other: &Node) -> Option<Ordering> {
23        Some(self.cmp(other))
24    }
25}
26impl Ord for Node {
27    fn cmp(&self, other: &Self) -> Ordering {
28        other.cost.cmp(&self.cost)
29    }
30}
31
32pub fn dijkstra(edge: &[Vec<(usize, i64)>], start: usize, end: usize) -> Option<i64> {
33    assert_ne!(start, end);
34    let mut costs = vec![None; edge.len()];
35    let mut nodes = BinaryHeap::new();
36    nodes.push(Node::new(start, 0));
37
38    while let Some(Node { position, cost }) = nodes.pop() {
39        if costs[position].is_some() {
40            continue;
41        }
42        if position == end {
43            return Some(cost);
44        }
45        costs[position] = Some(cost);
46
47        edge[position]
48            .iter()
49            .filter(|(to, c)| costs[*to].filter(|&d| d <= cost + c).is_none())
50            .for_each(|&(to, c)| {
51                nodes.push(Node::new(to, cost + c));
52            });
53    }
54    None
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    #[test]
61    fn test_dijkstra() {
62        let graph = vec![
63            vec![(2, 10), (1, 1)],
64            vec![(3, 2)],
65            vec![(1, 1), (3, 3), (4, 1)],
66            vec![(0, 7), (4, 2)],
67            vec![],
68        ];
69
70        assert_eq!(dijkstra(&graph, 0, 1), Some(1));
71        assert_eq!(dijkstra(&graph, 0, 2), Some(10));
72        assert_eq!(dijkstra(&graph, 0, 3), Some(3));
73        assert_eq!(dijkstra(&graph, 0, 4), Some(5));
74        assert_eq!(dijkstra(&graph, 3, 0), Some(7));
75        assert_eq!(dijkstra(&graph, 4, 0), None);
76    }
77
78    #[test]
79    fn chooses_the_shorter_indirect_path() {
80        let graph = vec![vec![(1, 10), (2, 1)], vec![], vec![(1, 1)]];
81        assert_eq!(dijkstra(&graph, 0, 1), Some(2));
82    }
83
84    #[test]
85    #[should_panic]
86    fn test_panic() {
87        dijkstra(&[], 0, 0);
88    }
89}