extra functionality for dijkstra (#66)

* extra functionality for dijkstra

* add example for looping over shortest path

* fix

---------

Co-authored-by: Luke Videckis <lukevideckis@gmail.com>
Co-authored-by: Cameron Custer <73217097+cameroncuster@users.noreply.github.com>

Author: lrvideckis

examples/graphs/dijk_aizu.rs

@@ -16,7 +16,10 @@ fn main() {
         adj[u].push((v, w));
     }
 
-    let dist = dijk(&adj, s);
+    let (dist, par) = dijk(&adj, s);
+
+    assert_eq!(par[s], None);
+
     for d in dist {
         if d == u64::MAX {
             println!("INF");

examples/graphs/dijk_yosupo.rs

@@ -17,37 +17,22 @@ fn main() {
         adj[u].push((v, w));
     }
 
-    let dist = dijk(&adj, s);
+    let (dist, par) = dijk(&adj, s);
+
+    assert_eq!(par[s], None);
 
     if dist[t] == u64::MAX {
         println!("{}", -1);
         return;
     }
 
-    let mut par = vec![0; n];
-    par[s] = s;
-    let mut seen = vec![false; n];
-    seen[s] = true;
-    let mut q = std::collections::VecDeque::new();
-    q.push_back(s);
-    while let Some(u) = q.pop_front() {
-        for &(v, w) in &adj[u] {
-            if seen[v] || dist[u] + w != dist[v] {
-                continue;
-            }
-            par[v] = u;
-            seen[v] = true;
-            q.push_back(v);
-        }
-    }
-
     let mut path = vec![];
     let mut u = t;
-    while u != s {
-        path.push(u);
-        u = par[u];
+    path.push(u);
+    while let Some(prev_u) = par[u] {
+        path.push(prev_u);
+        u = prev_u;
     }
-    path.push(s);
 
     path.reverse();
 

src/graphs/dijk.rs

@@ -4,40 +4,41 @@
 /// ```
 /// use programming_team_code_rust::graphs::dijk::dijk;
 ///
-/// let adj = vec![
-///    vec![(1, 1), (2, 4)],
-///    vec![(2, 2), (3, 5)],
-///    vec![(3, 1)],
-///    vec![],
-///    vec![(0, 7)],
-///    vec![(1, 3), (3, 2)],
-///    vec![(0, 2), (2, 3)],
-///    vec![(1, 2), (4, 2)],
-/// ];
+/// let mut adj = vec![vec![]; 4];
+/// for (u, v, w) in [(0, 1, 10), (1, 2, 100), (0, 2, 111), (2, 2, 0)] {
+///    adj[u].push((v, w));
+/// }
 ///
-/// let dist = dijk(&adj, 0);
-/// assert_eq!(dist, vec![0, 1, 3, 4, u64::MAX, u64::MAX, u64::MAX, u64::MAX]);
+/// let (dist, par) = dijk(&adj, 0);
+/// assert_eq!(dist, [0, 10, 110, u64::MAX]);
+/// assert_eq!(par, [None, Some(0), Some(1), None]);
+///
+/// let mut u = 2; // loop over shortest path in reverse
+/// while let Some(prev) = par[u] {
+///    u = prev;
+/// }
 /// ```
 ///
 /// # Complexity
 /// - V: number of vertices
 /// - E: number of edges
 /// - Time: O(V + E log E)
 /// - Space: O(V + E)
-pub fn dijk(adj: &[Vec<(usize, u64)>], s: usize) -> Vec<u64> {
+pub fn dijk(adj: &[Vec<(usize, u64)>], s: usize) -> (Vec<u64>, Vec<Option<usize>>) {
     use std::cmp::Reverse;
     let n = adj.len();
     let mut dist = vec![u64::MAX; n];
+    let mut par = vec![None; n];
     let mut q = std::collections::BinaryHeap::new();
-    q.push(Reverse((0, s)));
-    while let Some(Reverse((d, u))) = q.pop() {
+    q.push(Reverse((0, s, None)));
+    while let Some(Reverse((d, u, p))) = q.pop() {
         if dist[u] <= d {
             continue;
         }
-        dist[u] = d;
+        (dist[u], par[u]) = (d, p);
         for &(v, w) in &adj[u] {
-            q.push(Reverse((dist[u] + w, v)));
+            q.push(Reverse((dist[u] + w, v, Some(u))));
         }
     }
-    dist
+    (dist, par)
 }