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
 98
 99
100
101

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
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[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);
    }
}