From 49ebcd1e898280f9d51d8eddf04f7bb62d3456d0 Mon Sep 17 00:00:00 2001 From: galleonli Date: Mon, 2 Nov 2026 03:54:19 -0500 Subject: [PATCH] Add a distance cutoff to Dijkstra path lengths --- ...stra-distance-cutoff-9f3248a7d61c5e02.yaml | 12 ++ rustworkx/__init__.py | 8 +- rustworkx/__init__.pyi | 1 + rustworkx/rustworkx.pyi | 2 + src/shortest_path/mod.rs | 81 +++++++++- tests/test_dijkstra_cutoff.py | 141 ++++++++++++++++++ 6 files changed, 237 insertions(+), 8 deletions(-) create mode 100644 releasenotes/notes/dijkstra-distance-cutoff-9f3248a7d61c5e02.yaml create mode 100644 tests/test_dijkstra_cutoff.py diff --git a/releasenotes/notes/dijkstra-distance-cutoff-9f3248a7d61c5e02.yaml b/releasenotes/notes/dijkstra-distance-cutoff-9f3248a7d61c5e02.yaml new file mode 100644 index 0000000000..77d212df98 --- /dev/null +++ b/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. diff --git a/rustworkx/__init__.py b/rustworkx/__init__.py index 24660f76c5..09df02881d 100644 --- a/rustworkx/__init__.py +++ b/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") diff --git a/rustworkx/__init__.pyi b/rustworkx/__init__.pyi index bdd554d8cd..5192be812d 100644 --- a/rustworkx/__init__.pyi +++ b/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], diff --git a/rustworkx/rustworkx.pyi b/rustworkx/rustworkx.pyi index 0df0ee667d..e5dc6cef85 100644 --- a/rustworkx/rustworkx.pyi +++ b/rustworkx/rustworkx.pyi @@ -905,6 +905,7 @@ 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], @@ -912,6 +913,7 @@ def graph_dijkstra_shortest_path_lengths( 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], diff --git a/src/shortest_path/mod.rs b/src/shortest_path/mod.rs index 40b06c9103..59ada6ea82 100644 --- a/src/shortest_path/mod.rs +++ b/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( + graph: G, + start: NodeIndex, + goal: Option, + edge_cost: F, + cutoff: f64, +) -> PyResult +where + G: IntoEdges + Visitable, + F: FnMut(G::EdgeRef) -> PyResult, +{ + 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, goal: Option, + cutoff: Option, ) -> PyResult { 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> = 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, goal: Option, + cutoff: Option, ) -> PyResult { 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> = dijkstra( &graph.graph, start, diff --git a/tests/test_dijkstra_cutoff.py b/tests/test_dijkstra_cutoff.py new file mode 100644 index 0000000000..af1b9b7109 --- /dev/null +++ b/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)