Add a distance cutoff to Dijkstra path lengths

releasenotes/notes/dijkstra-distance-cutoff-9f3248a7d61c5e02.yaml

@@ -0,0 +1,12 @@
+features:
+  - |
+    Added an optional ``cutoff`` argument to
+    :func:`~rustworkx.dijkstra_shortest_path_lengths`,
+    :func:`~rustworkx.graph_dijkstra_shortest_path_lengths`, and
+    :func:`~rustworkx.digraph_dijkstra_shortest_path_lengths`.
+    The search only explores paths whose total weight is at most ``cutoff``,
+    avoiding traversal of distant parts of the graph. For example,
+    ``rx.dijkstra_shortest_path_lengths(graph, source, float, cutoff=10)``
+    returns distances up to and including 10. If a specified ``goal`` is beyond
+    the cutoff, the result is empty. Omitting ``cutoff`` keeps the existing
+    unbounded search. Cutoffs must be non-negative and not NaN.

rustworkx/__init__.py

@@ -564,7 +564,7 @@ def all_pairs_dijkstra_path_lengths(graph, edge_cost_fn):
 
 
 @_rustworkx_dispatch
-def dijkstra_shortest_path_lengths(graph, node, edge_cost_fn, goal=None):
+def dijkstra_shortest_path_lengths(graph, node, edge_cost_fn, goal=None, cutoff=None):
     """Compute the lengths of the shortest paths for a graph object using
     Dijkstra's algorithm.
 
@@ -579,11 +579,17 @@ def dijkstra_shortest_path_lengths(graph, node, edge_cost_fn, goal=None):
         When specified the traversal will stop when the goal is reached and
         the output dictionary will only have a single entry with the length
         of the shortest path to the goal node.
+    :param float cutoff: An optional maximum path cost, inclusive. Paths above
+        this cost are not explored. Must be non-negative and not NaN. If None,
+        the search is unbounded. Edge costs are evaluated only for edges reached
+        by the bounded search.
 
     :returns: A dictionary of the shortest paths from the provided node where
         the key is the node index of the end of the path and the value is the
         cost/sum of the weights of path
     :rtype: dict
+    :raises ValueError: If cutoff is negative or NaN, or an evaluated edge cost
+        is negative or NaN.
     """
     raise TypeError(f"Invalid Input Type {type(graph)} for graph")
 

rustworkx/__init__.pyi

@@ -407,6 +407,7 @@ def dijkstra_shortest_path_lengths(
     node: int,
     edge_cost_fn: Callable[[_T], float] | None,
     goal: int | None = ...,
+    cutoff: float | None = ...,
 ) -> PathLengthMapping: ...
 def k_shortest_path_lengths(
     graph: PyGraph[_S, _T] | PyDiGraph[_S, _T],

rustworkx/rustworkx.pyi

@@ -905,13 +905,15 @@ def digraph_dijkstra_shortest_path_lengths(
     edge_cost_fn: Callable[[_T], float] | None,
     /,
     goal: int | None = ...,
+    cutoff: float | None = ...,
 ) -> PathLengthMapping: ...
 def graph_dijkstra_shortest_path_lengths(
     graph: PyGraph[_S, _T],
     node: int,
     edge_cost_fn: Callable[[_T], float] | None,
     /,
     goal: int | None = ...,
+    cutoff: float | None = ...,
 ) -> PathLengthMapping: ...
 def digraph_all_pairs_bellman_ford_path_lengths(
     graph: PyDiGraph[_S, _T],

src/shortest_path/mod.rs

@@ -25,7 +25,7 @@ use numpy::{IntoPyArray, PyArray2};
 use petgraph::graph::NodeIndex;
 use petgraph::prelude::*;
 use petgraph::stable_graph::EdgeIndex;
-use petgraph::visit::NodeCount;
+use petgraph::visit::{Control, IntoEdges, NodeCount, Visitable};
 use pyo3::Python;
 use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError};
 use pyo3::prelude::*;
@@ -35,6 +35,7 @@ use rustworkx_core::shortest_path::{
     all_shortest_paths, astar, bellman_ford, dijkstra, k_shortest_path, negative_cycle_finder,
     single_source_all_shortest_paths,
 };
+use rustworkx_core::traversal::{DijkstraEvent, dijkstra_search};
 
 use crate::iterators::{
     AllPairsPathLengthMapping, AllPairsPathMapping, NodeIndices, NodesCountMapping,
@@ -544,6 +545,44 @@ pub fn digraph_has_path(
     Ok(!path_mapping.paths.is_empty())
 }
 
+fn bounded_dijkstra_lengths<G, F>(
+    graph: G,
+    start: NodeIndex,
+    goal: Option<NodeIndex>,
+    edge_cost: F,
+    cutoff: f64,
+) -> PyResult<PathLengthMapping>
+where
+    G: IntoEdges<NodeId = NodeIndex> + Visitable,
+    F: FnMut(G::EdgeRef) -> PyResult<f64>,
+{
+    if cutoff.is_nan() || cutoff < 0.0 {
+        return Err(PyValueError::new_err(
+            "cutoff must be non-negative and not NaN",
+        ));
+    }
+
+    // A bounded query may reach very few nodes, so avoid a graph-sized distance array.
+    let mut path_lengths = DictMap::new();
+    dijkstra_search(graph, Some(start), edge_cost, |event| {
+        if let DijkstraEvent::Discover(node, cost) = event {
+            if cost > cutoff {
+                return Control::Break(());
+            }
+            if let Some(goal) = goal {
+                if node == goal {
+                    path_lengths.insert(node.index(), cost);
+                    return Control::Break(());
+                }
+            } else if node != start {
+                path_lengths.insert(node.index(), cost);
+            }
+        }
+        Control::Continue
+    })?;
+    Ok(PathLengthMapping { path_lengths })
+}
+
 /// Compute the lengths of the shortest paths for a PyGraph object using
 /// Dijkstra's algorithm
 ///
@@ -557,21 +596,25 @@ pub fn digraph_has_path(
 ///     When specified the traversal will stop when the goal is reached and
 ///     the output dictionary will only have a single entry with the length
 ///     of the shortest path to the goal node.
+/// :param float cutoff: An optional maximum path cost, inclusive. Paths above
+///     this cost are not explored. Must be non-negative and not NaN. If None,
+///     the search is unbounded. Edge costs are evaluated only for edges reached
+///     by the bounded search.
 ///
 /// :returns: A dictionary of the shortest paths from the provided node where
 ///     the key is the node index of the end of the path and the value is the
 ///     cost/sum of the weights of path
 /// :rtype: PathLengthMapping
-/// :raises ValueError: when an edge weight with NaN or negative value
-///     is provided.
+/// :raises ValueError: when cutoff or an evaluated edge weight is NaN or negative.
 #[pyfunction]
-#[pyo3(text_signature = "(graph, node, edge_cost_fn, /, goal=None)", signature = (graph, node, edge_cost_fn, goal=None))]
+#[pyo3(text_signature = "(graph, node, edge_cost_fn, /, goal=None, cutoff=None)", signature = (graph, node, edge_cost_fn, goal=None, cutoff=None))]
 pub fn graph_dijkstra_shortest_path_lengths(
     py: Python,
     graph: &graph::PyGraph,
     node: usize,
     edge_cost_fn: Py<PyAny>,
     goal: Option<usize>,
+    cutoff: Option<f64>,
 ) -> PyResult<PathLengthMapping> {
     let start = NodeIndex::new(node);
     let edge_cost_callable = CostFn::from(edge_cost_fn);
@@ -583,6 +626,16 @@ pub fn graph_dijkstra_shortest_path_lengths(
         )));
     }
 
+    if let Some(cutoff) = cutoff {
+        return bounded_dijkstra_lengths(
+            &graph.graph,
+            start,
+            goal_index,
+            |e| edge_cost_callable.call(py, e.weight()),
+            cutoff,
+        );
+    }
+
     let res: Vec<Option<f64>> = dijkstra(
         &graph.graph,
         start,
@@ -632,21 +685,25 @@ pub fn graph_dijkstra_shortest_path_lengths(
 ///     When specified the traversal will stop when the goal is reached and
 ///     the output dictionary will only have a single entry with the length
 ///     of the shortest path to the goal node.
+/// :param float cutoff: An optional maximum path cost, inclusive. Paths above
+///     this cost are not explored. Must be non-negative and not NaN. If None,
+///     the search is unbounded. Edge costs are evaluated only for edges reached
+///     by the bounded search.
 ///
 /// :returns: A dictionary of the shortest paths from the provided node where
 ///     the key is the node index of the end of the path and the value is the
 ///     cost/sum of the weights of path
 /// :rtype: PathLengthMapping
-/// :raises ValueError: when an edge weight with NaN or negative value
-///     is provided.
+/// :raises ValueError: when cutoff or an evaluated edge weight is NaN or negative.
 #[pyfunction]
-#[pyo3(text_signature = "(graph, node, edge_cost_fn, /, goal=None)", signature = (graph, node, edge_cost_fn, goal=None))]
+#[pyo3(text_signature = "(graph, node, edge_cost_fn, /, goal=None, cutoff=None)", signature = (graph, node, edge_cost_fn, goal=None, cutoff=None))]
 pub fn digraph_dijkstra_shortest_path_lengths(
     py: Python,
     graph: &digraph::PyDiGraph,
     node: usize,
     edge_cost_fn: Py<PyAny>,
     goal: Option<usize>,
+    cutoff: Option<f64>,
 ) -> PyResult<PathLengthMapping> {
     let edge_cost_callable = CostFn::from(edge_cost_fn);
 
@@ -659,6 +716,16 @@ pub fn digraph_dijkstra_shortest_path_lengths(
         )));
     }
 
+    if let Some(cutoff) = cutoff {
+        return bounded_dijkstra_lengths(
+            &graph.graph,
+            start,
+            goal_index,
+            |e| edge_cost_callable.call(py, e.weight()),
+            cutoff,
+        );
+    }
+
     let res: Vec<Option<f64>> = dijkstra(
         &graph.graph,
         start,

tests/test_dijkstra_cutoff.py

@@ -0,0 +1,141 @@
+# 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.
+
+import unittest
+
+import rustworkx
+
+
+class TestDijkstraCutoff(unittest.TestCase):
+    graph_types = (rustworkx.PyGraph, rustworkx.PyDiGraph)
+
+    def test_inclusive_boundary_and_shorter_alternative(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            graph.add_nodes_from(range(4))
+            graph.add_edges_from([(0, 1, 20), (0, 1, 2), (0, 2, 10), (1, 2, 3), (2, 3, 1)])
+            with self.subTest(graph_type=graph_type):
+                self.assertEqual(
+                    {1: 2.0, 2: 5.0},
+                    rustworkx.dijkstra_shortest_path_lengths(graph, 0, float, cutoff=5.0),
+                )
+
+    def test_zero_weight_paths_and_cycles(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            graph.add_nodes_from(range(4))
+            graph.add_edges_from([(0, 1, 0), (1, 0, 0), (1, 2, 0), (2, 3, 1)])
+            with self.subTest(graph_type=graph_type):
+                self.assertEqual(
+                    {1: 0.0, 2: 0.0},
+                    rustworkx.dijkstra_shortest_path_lengths(graph, 0, float, cutoff=0.0),
+                )
+
+    def test_goal_inside_and_outside_cutoff(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            graph.add_nodes_from(range(4))
+            graph.add_edges_from([(0, 1, 2), (1, 2, 3)])
+            for goal, expected in [(0, {0: 0.0}), (1, {1: 2.0}), (2, {}), (3, {})]:
+                with self.subTest(graph_type=graph_type, goal=goal):
+                    self.assertEqual(
+                        expected,
+                        rustworkx.dijkstra_shortest_path_lengths(graph, 0, float, goal, cutoff=2.0),
+                    )
+
+    def test_none_and_infinity_match_unbounded_search(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            graph.add_nodes_from(range(3))
+            graph.add_edges_from([(0, 1, 2), (1, 2, float("inf"))])
+            expected = rustworkx.dijkstra_shortest_path_lengths(graph, 0, float)
+            for cutoff in [None, float("inf")]:
+                with self.subTest(graph_type=graph_type, cutoff=cutoff):
+                    self.assertEqual(
+                        expected,
+                        rustworkx.dijkstra_shortest_path_lengths(graph, 0, float, cutoff=cutoff),
+                    )
+
+    def test_invalid_cutoff(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            graph.add_node(None)
+            for cutoff in [-1.0, float("-inf"), float("nan")]:
+                with self.subTest(graph_type=graph_type, cutoff=cutoff):
+                    with self.assertRaisesRegex(ValueError, "cutoff"):
+                        rustworkx.dijkstra_shortest_path_lengths(graph, 0, float, cutoff=cutoff)
+
+    def test_direct_graph_functions(self):
+        functions = (
+            (rustworkx.PyGraph, rustworkx.graph_dijkstra_shortest_path_lengths),
+            (rustworkx.PyDiGraph, rustworkx.digraph_dijkstra_shortest_path_lengths),
+        )
+        for graph_type, function in functions:
+            graph = graph_type()
+            graph.add_nodes_from(range(3))
+            graph.add_edges_from([(0, 1, 1), (1, 2, 1)])
+            with self.subTest(function=function):
+                self.assertEqual({1: 1.0}, function(graph, 0, float, cutoff=1.0))
+
+    def test_stops_before_visiting_distant_edges(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            graph.add_nodes_from(range(100))
+            graph.add_edges_from((i, i + 1, i) for i in range(99))
+            seen = []
+
+            def cost(edge):
+                seen.append(edge)
+                return 1.0
+
+            with self.subTest(graph_type=graph_type):
+                result = rustworkx.dijkstra_shortest_path_lengths(graph, 0, cost, cutoff=2.0)
+                self.assertEqual({1: 1.0, 2: 2.0}, result)
+                self.assertEqual([0, 1, 2], seen)
+
+    def test_callback_errors_and_invalid_weights(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            graph.add_nodes_from(range(2))
+            graph.add_edge(0, 1, None)
+            for weight in [-1.0, float("nan")]:
+                with self.subTest(graph_type=graph_type, weight=weight):
+                    with self.assertRaises(ValueError):
+                        rustworkx.dijkstra_shortest_path_lengths(
+                            graph, 0, lambda _: weight, cutoff=0.0
+                        )
+
+            def cost(_):
+                raise RuntimeError("edge callback failed")
+
+            with self.subTest(graph_type=graph_type):
+                with self.assertRaisesRegex(RuntimeError, "edge callback failed"):
+                    rustworkx.dijkstra_shortest_path_lengths(graph, 0, cost, cutoff=0.0)
+
+    def test_removed_nodes_and_self_loops(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            graph.add_nodes_from(range(5))
+            graph.remove_node(1)
+            graph.add_edges_from([(0, 0, 0), (0, 2, 1), (2, 3, 1), (3, 4, 1)])
+            with self.subTest(graph_type=graph_type):
+                self.assertEqual(
+                    {2: 1.0, 3: 2.0},
+                    rustworkx.dijkstra_shortest_path_lengths(graph, 0, float, cutoff=2.0),
+                )
+
+    def test_invalid_source(self):
+        for graph_type in self.graph_types:
+            graph = graph_type()
+            with self.subTest(graph_type=graph_type):
+                with self.assertRaises(IndexError):
+                    rustworkx.dijkstra_shortest_path_lengths(graph, 0, float, cutoff=0.0)