Use Vector for intermediate computations in Dijkstra

retworkx-core/src/distancemap.rs

@@ -0,0 +1,88 @@
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations
+// under the License.
+
+//! This module contains the [`DistanceMap`] trait which is used in
+//! [`shortest_path`](crate::shortest_path).
+//!
+//! The trait allows the shortest path functions to support multiple
+//! return types.
+
+use std::hash::Hash;
+
+use petgraph::graph::IndexType;
+
+use crate::dictmap::*;
+use hashbrown::HashMap;
+
+/// A mapping for storing the distances of nodes for shortest path algorithms.
+pub trait DistanceMap<K, V> {
+    /// Create mapping with support for items between 0 and `num_elements - 1`.
+    fn build(num_elements: usize) -> Self;
+
+    /// Get the distance to the item at `pos`. If the distance does not exist,
+    /// the function returns `None`.
+    fn get_item(&self, pos: K) -> Option<&V>;
+
+    /// Insert item at position `pos` with distance `V`.
+    fn put_item(&mut self, pos: K, val: V);
+}
+
+impl<K: IndexType, V: Clone> DistanceMap<K, V> for Vec<Option<V>> {
+    #[inline]
+    fn build(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 build(_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);
+    }
+}
+
+impl<K: Eq + Hash, V: Clone> DistanceMap<K, V> for HashMap<K, V> {
+    #[inline]
+    fn build(_num_elements: usize) -> Self {
+        HashMap::<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);
+    }
+}

retworkx-core/src/lib.rs

@@ -76,6 +76,7 @@ pub mod shortest_path;
 pub mod traversal;
 // These modules define additional data structures
 pub mod dictmap;
+pub mod distancemap;
 mod min_scored;
 
 // re-export petgraph so there is a consistent version available to users and

retworkx-core/src/shortest_path/dijkstra.rs

@@ -20,12 +20,11 @@
 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::visit::{EdgeRef, IntoEdges, NodeIndexable, VisitMap, Visitable};
 
 use crate::dictmap::*;
+use crate::distancemap::DistanceMap;
 use crate::min_scored::MinScored;
 
 /// Dijkstra's shortest path algorithm.
@@ -45,7 +44,7 @@ use crate::min_scored::MinScored;
 /// the value is a Vec of node indices of the path starting with `start` and
 /// ending at the index.
 ///
-/// Returns a [`DictMap`] that maps `NodeId` to path cost.
+/// Returns a [`DistanceMap`] that maps `NodeId` to path cost.
 /// # Example
 /// ```rust
 /// use retworkx_core::petgraph::Graph;
@@ -98,24 +97,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: IntoEdges + Visitable + NodeIndexable,
     G::NodeId: Eq + Hash,
     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::build(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 +134,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 +155,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 +169,6 @@ where
         }
         visited.visit(node);
     }
+
     Ok(scores)
 }

retworkx-core/src/shortest_path/k_shortest_path.rs

@@ -22,10 +22,11 @@ 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::*;
+use crate::distancemap::DistanceMap;
 use crate::min_scored::MinScored;
 
 /// k'th shortest path algorithm.
@@ -42,28 +43,28 @@ use crate::min_scored::MinScored;
 ///
 /// Computes in **O(k * (|E| + |V|*log(|V|)))** time (average).
 ///
-/// Returns a [`DictMap`] that maps `NodeId` to path cost as the value.
+/// Returns a [`DistanceMap`] that maps `NodeId` to path cost as the value.
 ///
 /// # Example:
 /// ```rust
 ///
 /// use retworkx_core::petgraph;
 /// use retworkx_core::petgraph::graph::NodeIndex;
 /// use retworkx_core::shortest_path::k_shortest_path;
-/// use retworkx_core::dictmap::DictMap;
+/// use hashbrown::HashMap;
 /// use retworkx_core::Result;
 ///
 /// let g = petgraph::graph::UnGraph::<i32, _>::from_edges(&[
 ///     (0, 1), (1, 2), (2, 3), (3, 0), (4, 5), (1, 4), (5, 6), (6, 7), (7, 5)
 /// ]);
 ///
-/// let res: Result<DictMap<NodeIndex, f64>> = k_shortest_path(
+/// let res: Result<HashMap<NodeIndex, f64>> = k_shortest_path(
 ///     &g, NodeIndex::new(1), None, 2,
 ///     |e: retworkx_core::petgraph::graph::EdgeReference<&'static str>| Ok(1.0),
 /// );
 ///
 /// let output = res.unwrap();
-/// let expected: DictMap<NodeIndex, f64> = [
+/// let expected: HashMap<NodeIndex, f64> = [
 ///     (NodeIndex::new(0), 3.0),
 ///     (NodeIndex::new(1), 2.0),
 ///     (NodeIndex::new(2), 3.0),
@@ -75,21 +76,22 @@ use crate::min_scored::MinScored;
 /// ].iter().cloned().collect();
 /// assert_eq!(expected, output);
 /// ```
-pub fn k_shortest_path<G, F, E, K>(
+pub fn k_shortest_path<G, F, E, K, S>(
     graph: G,
     start: G::NodeId,
     goal: Option<G::NodeId>,
     k: usize,
     mut edge_cost: F,
-) -> Result<DictMap<G::NodeId, K>, E>
+) -> Result<S, 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,
+    S: DistanceMap<G::NodeId, K>,
 {
     let mut counter: Vec<usize> = vec![0; graph.node_bound()];
-    let mut scores = DictMap::with_capacity(graph.node_count());
+    let mut scores: S = S::build(graph.node_bound());
     let mut visit_next = BinaryHeap::new();
     let zero_score = K::default();
 
@@ -104,7 +106,7 @@ where
         }
 
         if current_counter == k {
-            scores.insert(node, node_score);
+            scores.put_item(node, node_score);
         }
 
         //Already reached goal k times
@@ -117,5 +119,6 @@ where
                 .push(MinScored(node_score + edge_cost(edge)?, edge.target()));
         }
     }
+
     Ok(scores)
 }

src/shortest_path/all_pairs_dijkstra.rs

@@ -80,24 +80,21 @@ pub fn all_pairs_dijkstra_path_lengths<Ty: EdgeType + Sync>(
     let out_map: DictMap<usize, PathLengthMapping> = node_indices
         .into_par_iter()
         .map(|x| {
+            let path_lenghts: PyResult<Vec<Option<f64>>> =
+                dijkstra(graph, x, None, |e| edge_cost(e.id()), None);
             let out_map = PathLengthMapping {
-                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(),
+                path_lengths: path_lenghts
+                    .unwrap()
+                    .into_iter()
+                    .enumerate()
+                    .filter_map(|(index, opt_cost)| {
+                        if index != x.index() {
+                            opt_cost.map(|cost| (index, cost))
+                        } else {
+                            None
+                        }
+                    })
+                    .collect(),
             };
             (x.index(), out_map)
         })