Tentative all_shortest_paths scaffolding

Qiskit · Oct 23, 2023 · 68ada54 · 68ada54
1 parent 4018b35
commit 68ada54
Showing 1 changed file with 70 additions and 6 deletions.
diff --git a/rustworkx-core/src/shortest_path/dijkstra.rs b/rustworkx-core/src/shortest_path/dijkstra.rs
@@ -102,7 +102,7 @@ pub fn dijkstra<G, F, K, E, S>(
     start: G::NodeId,
     goal: Option<G::NodeId>,
     mut edge_cost: F,
-    mut path: Option<&mut DictMap<G::NodeId, Vec<G::NodeId>>>,
+    mut path: Option<&mut DictMap<G::NodeId, Vec<Vec<G::NodeId>>>>,
 ) -> Result<S, E>
 where
     G: IntoEdges + Visitable + NodeIndexable,
@@ -118,7 +118,7 @@ where
     scores.put_item(start, zero_score);
     visit_next.push(MinScored(zero_score, start));
     if path.is_some() {
-        path.as_mut().unwrap().insert(start, vec![start]);
+        path.as_mut().unwrap().insert(start, vec![vec![start]]);
     }
     while let Some(MinScored(node_score, node)) = visit_next.pop() {
         if visited.is_visited(&node) {
@@ -140,10 +140,34 @@ where
                         scores.put_item(next, next_score);
                         visit_next.push(MinScored(next_score, next));
                         if path.is_some() {
-                            let mut node_path = path.as_mut().unwrap().get(&node).unwrap().clone();
+                            let mut node_path = path
+                                .as_mut()
+                                .unwrap()
+                                .get(&node)
+                                .unwrap()
+                                .iter()
+                                .find(|x| x.last().unwrap() == &node)
+                                .unwrap()
+                                .clone();
                             node_path.push(next);
                             path.as_mut().unwrap().entry(next).and_modify(|new_vec| {
-                                *new_vec = node_path;
+                                *new_vec = vec![node_path];
+                            });
+                        }
+                    } else if next_score == *current_score {
+                        if path.is_some() {
+                            let mut node_path = path
+                                .as_mut()
+                                .unwrap()
+                                .get(&node)
+                                .unwrap()
+                                .iter()
+                                .find(|x| x.last().unwrap() == &node)
+                                .unwrap()
+                                .clone();
+                            node_path.push(next);
+                            path.as_mut().unwrap().entry(next).and_modify(|new_vec| {
+                                new_vec.push(node_path);
                             });
                         }
                     }
@@ -152,9 +176,20 @@ where
                     scores.put_item(next, next_score);
                     visit_next.push(MinScored(next_score, next));
                     if path.is_some() {
-                        let mut node_path = path.as_mut().unwrap().get(&node).unwrap().clone();
+                        let mut node_path = path
+                            .as_mut()
+                            .unwrap()
+                            .get(&node)
+                            .unwrap()
+                            .iter()
+                            .find(|x| x.last().unwrap() == &node)
+                            .unwrap()
+                            .clone();
                         node_path.push(next);
-                        path.as_mut().unwrap().entry(next).or_insert(node_path);
+                        path.as_mut()
+                            .unwrap()
+                            .entry(next)
+                            .or_insert(vec![node_path]);
                     }
                 }
             }
@@ -168,6 +203,7 @@ where
 #[cfg(test)]
 mod tests {
     use crate::dictmap::DictMap;
+    use crate::dictmap::InitWithHasher;
     use crate::shortest_path::dijkstra;
     use crate::Result;
     use petgraph::prelude::*;
@@ -226,4 +262,32 @@ mod tests {
             dijkstra(&g, a, Some(c), |e| Ok(*e.weight()), None);
         assert_eq!(scores.unwrap()[&c], 9);
     }
+
+    #[test]
+    fn test_dijk_all_paths() {
+        let mut g = Graph::new_undirected();
+        let a = g.add_node("A");
+        let b = g.add_node("B");
+        let c = g.add_node("C");
+        let d = g.add_node("D");
+        let e = g.add_node("E");
+        let f = g.add_node("F");
+        g.add_edge(a, b, 1);
+        g.add_edge(b, c, 1);
+        g.add_edge(a, d, 1);
+        g.add_edge(d, c, 1);
+        g.add_edge(a, e, 1);
+        g.add_edge(e, f, 1);
+        g.add_edge(f, c, 1);
+        g.add_edge(e, c, 1);
+        println!("{:?}", g);
+
+        let mut paths: DictMap<NodeIndex, Vec<Vec<NodeIndex>>> =
+            DictMap::with_capacity(g.node_count());
+
+        let scores: Result<DictMap<NodeIndex, usize>> =
+            dijkstra(&g, a, Some(c), |e| Ok(*e.weight()), Some(&mut paths));
+        assert_eq!(scores.unwrap()[&c], 2);
+        assert_eq!(paths.get(&c).unwrap().len(), 3);
+    }
 }