Use Vector for intermediate computations in Dijkstra

retworkx-core/src/shortest_path/dijkstra.rs

@@ -20,14 +20,54 @@
 use std::collections::BinaryHeap;
 use std::hash::Hash;
 
-use indexmap::map::Entry::{Occupied, Vacant};
-
 use petgraph::algo::Measure;
-use petgraph::visit::{EdgeRef, IntoEdges, VisitMap, Visitable};
+use petgraph::graph::IndexType;
+use petgraph::visit::{EdgeRef, IntoEdges, NodeIndexable, VisitMap, Visitable};
 
 use crate::dictmap::*;
 use crate::min_scored::MinScored;
 
+// Trait for supporting storing
+pub trait DistanceMap<K, V> {
+    fn new_distancemap(num_elements: usize) -> Self;
+    fn get_item(&self, pos: K) -> Option<&V>;
+    fn put_item(&mut self, pos: K, val: V);
+}
+
+impl<K: IndexType, V: Clone> DistanceMap<K, V> for Vec<Option<V>> {
+    #[inline]
+    fn new_distancemap(num_elements: usize) -> Self {
+        vec![None; num_elements]
+    }
+
+    #[inline]
+    fn get_item(&self, pos: K) -> Option<&V> {
+        self[pos.index()].as_ref()
+    }
+
+    #[inline]
+    fn put_item(&mut self, pos: K, val: V) {
+        self[pos.index()] = Some(val);
+    }
+}
+
+impl<K: Eq + Hash, V: Clone> DistanceMap<K, V> for DictMap<K, V> {
+    #[inline]
+    fn new_distancemap(_num_elements: usize) -> Self {
+        DictMap::<K, V>::default()
+    }
+
+    #[inline]
+    fn get_item(&self, pos: K) -> Option<&V> {
+        self.get(&pos)
+    }
+
+    #[inline]
+    fn put_item(&mut self, pos: K, val: V) {
+        self.insert(pos, val);
+    }
+}
+
 /// Dijkstra's shortest path algorithm.
 ///
 /// Compute the length of the shortest path from `start` to every reachable
@@ -98,24 +138,25 @@ use crate::min_scored::MinScored;
 /// assert_eq!(res.unwrap(), expected_res);
 /// // z is not inside res because there is not path from b to z.
 /// ```
-pub fn dijkstra<G, F, K, E>(
+pub fn dijkstra<G, F, K, E, S>(
     graph: G,
     start: G::NodeId,
     goal: Option<G::NodeId>,
     mut edge_cost: F,
     mut path: Option<&mut DictMap<G::NodeId, Vec<G::NodeId>>>,
-) -> Result<DictMap<G::NodeId, K>, E>
+) -> Result<S, E>
 where
-    G: IntoEdges + Visitable,
-    G::NodeId: Eq + Hash,
+    G: IntoEdges + Visitable + NodeIndexable,
+    G::NodeId: Eq + Hash + IndexType,
     F: FnMut(G::EdgeRef) -> Result<K, E>,
     K: Measure + Copy,
+    S: DistanceMap<G::NodeId, K>,
 {
     let mut visited = graph.visit_map();
-    let mut scores = DictMap::new();
+    let mut scores: S = S::new_distancemap(graph.node_bound());
     let mut visit_next = BinaryHeap::new();
     let zero_score = K::default();
-    scores.insert(start, zero_score);
+    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]);
@@ -134,10 +175,10 @@ where
             }
             let cost = edge_cost(edge)?;
             let next_score = node_score + cost;
-            match scores.entry(next) {
-                Occupied(ent) => {
-                    if next_score < *ent.get() {
-                        *ent.into_mut() = next_score;
+            match scores.get_item(next) {
+                Some(current_score) => {
+                    if next_score < *current_score {
+                        scores.put_item(next, next_score);
                         visit_next.push(MinScored(next_score, next));
                         if path.is_some() {
                             let mut node_path = path
@@ -155,8 +196,8 @@ where
                         }
                     }
                 }
-                Vacant(ent) => {
-                    ent.insert(next_score);
+                None => {
+                    scores.put_item(next, next_score);
                     visit_next.push(MinScored(next_score, next));
                     if path.is_some() {
                         let mut node_path =
@@ -169,5 +210,6 @@ where
         }
         visited.visit(node);
     }
+
     Ok(scores)
 }

retworkx-core/src/shortest_path/k_shortest_path.rs

@@ -22,7 +22,8 @@ use std::hash::Hash;
 
 use petgraph::algo::Measure;
 use petgraph::visit::{
-    EdgeRef, IntoEdges, NodeCount, NodeIndexable, Visitable,
+    EdgeRef, IntoEdges, IntoNodeIdentifiers, NodeCount, NodeIndexable,
+    Visitable,
 };
 
 use crate::dictmap::*;
@@ -83,13 +84,13 @@ pub fn k_shortest_path<G, F, E, K>(
     mut edge_cost: F,
 ) -> Result<DictMap<G::NodeId, K>, E>
 where
-    G: IntoEdges + Visitable + NodeCount + NodeIndexable,
+    G: IntoEdges + Visitable + NodeCount + NodeIndexable + IntoNodeIdentifiers,
     G::NodeId: Eq + Hash,
     F: FnMut(G::EdgeRef) -> Result<K, E>,
     K: Measure + Copy,
 {
     let mut counter: Vec<usize> = vec![0; graph.node_bound()];
-    let mut scores = DictMap::with_capacity(graph.node_count());
+    let mut scores: Vec<Option<K>> = vec![None; graph.node_bound()];
     let mut visit_next = BinaryHeap::new();
     let zero_score = K::default();
 
@@ -104,7 +105,7 @@ where
         }
 
         if current_counter == k {
-            scores.insert(node, node_score);
+            scores[graph.to_index(node)] = Some(node_score);
         }
 
         //Already reached goal k times
@@ -117,5 +118,12 @@ where
                 .push(MinScored(node_score + edge_cost(edge)?, edge.target()));
         }
     }
-    Ok(scores)
+
+    Ok(scores
+        .into_iter()
+        .enumerate()
+        .filter_map(|(node, score)| {
+            score.map(|val| (graph.from_index(node), val))
+        })
+        .collect())
 }

src/shortest_path/all_pairs_dijkstra.rs

@@ -81,23 +81,24 @@ pub fn all_pairs_dijkstra_path_lengths<Ty: EdgeType + Sync>(
         .into_par_iter()
         .map(|x| {
             let out_map = PathLengthMapping {
-                path_lengths: dijkstra(
+                path_lengths: (dijkstra(
                     graph,
                     x,
                     None,
                     |e| edge_cost(e.id()),
                     None,
                 )
-                .unwrap()
-                .iter()
-                .filter_map(|(index, cost)| {
-                    if *index == x {
-                        None
-                    } else {
-                        Some((index.index(), *cost))
-                    }
-                })
-                .collect(),
+                    as PyResult<DictMap<NodeIndex, f64>>)
+                    .unwrap()
+                    .iter()
+                    .filter_map(|(index, cost)| {
+                        if *index == x {
+                            None
+                        } else {
+                            Some((index.index(), *cost))
+                        }
+                    })
+                    .collect(),
             };
             (x.index(), out_map)
         })

src/shortest_path/mod.rs

@@ -75,13 +75,14 @@ pub fn graph_dijkstra_shortest_paths(
         DictMap::with_capacity(graph.node_count());
 
     let cost_fn = CostFn::try_from((weight_fn, default_weight))?;
-    dijkstra(
+
+    (dijkstra(
         &graph.graph,
         start,
         goal_index,
         |e| cost_fn.call(py, e.weight()),
         Some(&mut paths),
-    )?;
+    ) as PyResult<DictMap<NodeIndex, f64>>)?;
 
     Ok(PathMapping {
         paths: paths
@@ -144,7 +145,7 @@ pub fn digraph_dijkstra_shortest_paths(
     let cost_fn = CostFn::try_from((weight_fn, default_weight))?;
 
     if as_undirected {
-        dijkstra(
+        (dijkstra(
             // TODO: Use petgraph undirected adapter after
             // https://github.com/petgraph/petgraph/pull/318 is available in
             // a petgraph release.
@@ -153,15 +154,15 @@ pub fn digraph_dijkstra_shortest_paths(
             goal_index,
             |e| cost_fn.call(py, e.weight()),
             Some(&mut paths),
-        )?;
+        ) as PyResult<DictMap<NodeIndex, f64>>)?;
     } else {
-        dijkstra(
+        (dijkstra(
             &graph.graph,
             start,
             goal_index,
             |e| cost_fn.call(py, e.weight()),
             Some(&mut paths),
-        )?;
+        ) as PyResult<DictMap<NodeIndex, f64>>)?;
     }
     Ok(PathMapping {
         paths: paths
@@ -216,22 +217,36 @@ pub fn graph_dijkstra_shortest_path_lengths(
     let start = NodeIndex::new(node);
     let goal_index: Option<NodeIndex> = goal.map(NodeIndex::new);
 
-    let res = dijkstra(
+    let res: Vec<Option<f64>> = dijkstra(
         &graph.graph,
         start,
         goal_index,
         |e| edge_cost_callable.call(py, e.weight()),
         None,
     )?;
+
+    if let Some(goal_usize) = goal {
+        return Ok(PathLengthMapping {
+            path_lengths: match res[goal_usize] {
+                Some(goal_length) => {
+                    let mut ans = DictMap::new();
+                    ans.insert(goal_usize, goal_length);
+                    ans
+                }
+                None => DictMap::new(),
+            },
+        });
+    }
+
     Ok(PathLengthMapping {
         path_lengths: res
-            .iter()
-            .filter_map(|(k, v)| {
-                let k_int = k.index();
-                if k_int == node || goal.is_some() && goal.unwrap() != k_int {
-                    None
+            .into_iter()
+            .enumerate()
+            .filter_map(|(k_int, opt_v)| {
+                if k_int != node {
+                    opt_v.map(|v| (k_int, v))
                 } else {
-                    Some((k_int, *v))
+                    None
                 }
             })
             .collect(),
@@ -272,22 +287,36 @@ pub fn digraph_dijkstra_shortest_path_lengths(
     let start = NodeIndex::new(node);
     let goal_index: Option<NodeIndex> = goal.map(NodeIndex::new);
 
-    let res = dijkstra(
+    let res: Vec<Option<f64>> = dijkstra(
         &graph.graph,
         start,
         goal_index,
         |e| edge_cost_callable.call(py, e.weight()),
         None,
     )?;
+
+    if let Some(goal_usize) = goal {
+        return Ok(PathLengthMapping {
+            path_lengths: match res[goal_usize] {
+                Some(goal_length) => {
+                    let mut ans = DictMap::new();
+                    ans.insert(goal_usize, goal_length);
+                    ans
+                }
+                None => DictMap::new(),
+            },
+        });
+    }
+
     Ok(PathLengthMapping {
         path_lengths: res
-            .iter()
-            .filter_map(|(k, v)| {
-                let k_int = k.index();
-                if k_int == node || goal.is_some() && goal.unwrap() != k_int {
-                    None
+            .into_iter()
+            .enumerate()
+            .filter_map(|(k_int, opt_v)| {
+                if k_int != node {
+                    opt_v.map(|v| (k_int, v))
                 } else {
-                    Some((k_int, *v))
+                    None
                 }
             })
             .collect(),

src/steiner_tree.rs

@@ -162,7 +162,7 @@ fn fast_metric_edges(
     };
 
     let mut paths = DictMap::with_capacity(graph.graph.node_count());
-    let mut distance =
+    let mut distance: DictMap<NodeIndex, f64> =
         dijkstra(&graph.graph, dummy, None, cost_fn, Some(&mut paths))?;
     paths.remove(&dummy);
     distance.remove(&dummy);

tests/digraph/test_dijkstra.py

@@ -44,6 +44,16 @@ def test_dijkstra(self):
         expected = {4: 23.0}
         self.assertEqual(expected, path)
 
+    def test_dijkstra_length_with_no_path(self):
+        g = retworkx.PyDiGraph()
+        a = g.add_node("A")
+        b = g.add_node("B")
+        path_lenghts = retworkx.digraph_dijkstra_shortest_path_lengths(
+            g, a, edge_cost_fn=float, goal=b
+        )
+        expected = {}
+        self.assertEqual(expected, path_lenghts)
+
     def test_dijkstra_path(self):
         paths = retworkx.digraph_dijkstra_shortest_paths(self.graph, self.a)
         expected = {

tests/graph/test_dijkstra.py

@@ -57,6 +57,16 @@ def test_dijkstra_with_no_goal_set(self):
         expected = {1: 1.0, 2: 1.0, 3: 1.0, 4: 2.0, 5: 2.0}
         self.assertEqual(expected, path)
 
+    def test_dijkstra_length_with_no_path(self):
+        g = retworkx.PyGraph()
+        a = g.add_node("A")
+        b = g.add_node("B")
+        path_lenghts = retworkx.graph_dijkstra_shortest_path_lengths(
+            g, a, edge_cost_fn=float, goal=b
+        )
+        expected = {}
+        self.assertEqual(expected, path_lenghts)
+
     def test_dijkstra_path_with_no_goal_set(self):
         path = retworkx.graph_dijkstra_shortest_paths(self.graph, self.a)
         expected = {