1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

use std::cmp::Reverse;
use std::collections::BinaryHeap;

/// # Dijkstra
///
/// Example:
/// ```
/// use competitive_hpp::dijkstra::Dijkstra;
/// // edge: Vec<(from, to, cost)>
/// let edges: Vec<(usize, usize, i64)> = vec![(0, 1, 1),(0, 2, 6),(1, 3, 2)];
///
/// //Dijkstra::new(vertex num, edges, start vertex)
/// let dijkstra = Dijkstra::new(4, &edges, 0);
/// ```

#[derive(Clone, Debug)]
pub struct Dijkstra {
    dist: Vec<i64>,
    adjacency_list: Vec<Vec<(usize, i64)>>,
    n: usize,
}

impl Dijkstra {
    pub fn new(n: usize, edges: &[(usize, usize, i64)], start: usize) -> Self {
        let inf = i64::max_value();

        let mut dist: Vec<i64> = vec![inf; n];
        let adjacency_list = Self::create_adjacency_list(n, &edges);

        // MinHeap
        let mut heap: BinaryHeap<Reverse<(i64, usize)>> = BinaryHeap::new();

        dist[start] = 0;
        heap.push(Reverse((0, start)));

        while !heap.is_empty() {
            let Reverse((d, v)) = heap.pop().unwrap();

            if dist[v] < d {
                continue;
            }

            for &(u, cost) in &adjacency_list[v] {
                if dist[u] > dist[v] + cost {
                    dist[u] = dist[v] + cost;
                    heap.push(Reverse((dist[u], u)));
                }
            }
        }

        Dijkstra {
            dist,
            adjacency_list,
            n,
        }
    }

    fn create_adjacency_list(n: usize, edges: &[(usize, usize, i64)]) -> Vec<Vec<(usize, i64)>> {
        let mut adjacency_list: Vec<Vec<(usize, i64)>> = vec![vec![]; n];

        for &(from, to, cost) in edges {
            adjacency_list[from].push((to, cost))
        }

        adjacency_list
    }
}

#[test]
fn dijkstra_test() {
    // | -------2------->|
    // 0 --1--> 1 --2--> 3 --|
    // |---6--> 2 --2--->|   |
    // |<---------4----------|
    let edges = vec![
        (0, 1, 1),
        (0, 2, 6),
        (1, 3, 2),
        (2, 3, 2),
        (0, 3, 2),
        (3, 0, 4),
    ];

    let dijkstra = Dijkstra::new(4, &edges, 0);

    assert_eq!(dijkstra.dist[0], 0);
    assert_eq!(dijkstra.dist[1], 1);
    assert_eq!(dijkstra.dist[2], 6);
    assert_eq!(dijkstra.dist[3], 2);

    let dijkstra_another = Dijkstra::new(4, &edges, 1);

    assert_eq!(dijkstra_another.dist[0], 6);
    assert_eq!(dijkstra_another.dist[1], 0);
    assert_eq!(dijkstra_another.dist[2], 12);
    assert_eq!(dijkstra_another.dist[3], 2);
}