diff --git a/docs/source/api.rst b/docs/source/api.rst index 5e161e159..e381e5628 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -57,12 +57,14 @@ Traversal retworkx.dfs_edges retworkx.bfs_successors + retworkx.bfs_search retworkx.topological_sort retworkx.lexicographical_topological_sort retworkx.descendants retworkx.ancestors retworkx.collect_runs retworkx.collect_bicolor_runs + retworkx.visit.BFSVisitor .. _dag-algorithms: @@ -254,6 +256,7 @@ the functions from the explicitly typed based on the data type. retworkx.digraph_num_shortest_paths_unweighted retworkx.digraph_betweenness_centrality retworkx.digraph_unweighted_average_shortest_path_length + retworkx.digraph_bfs_search .. _api-functions-pygraph: @@ -296,6 +299,7 @@ typed API based on the data type. retworkx.graph_num_shortest_paths_unweighted retworkx.graph_betweenness_centrality retworkx.graph_unweighted_average_shortest_path_length + retworkx.graph_bfs_search Exceptions ========== @@ -310,6 +314,8 @@ Exceptions retworkx.NoSuitableNeighbors retworkx.NoPathFound retworkx.NullGraph + retworkx.visit.StopSearch + retworkx.visit.PruneSearch Custom Return Types =================== diff --git a/releasenotes/notes/bfs-search-da8ba99e8ecdd7ba.yaml b/releasenotes/notes/bfs-search-da8ba99e8ecdd7ba.yaml new file mode 100644 index 000000000..d3b486865 --- /dev/null +++ b/releasenotes/notes/bfs-search-da8ba99e8ecdd7ba.yaml @@ -0,0 +1,29 @@ +--- +features: + - | + Added a new :func:`~retworkx.bfs_search` (and it's per type variants + :func:`~retworkx.graph_bfs_search` and :func:`~retworkx.digraph_bfs_search`) + that traverses the graph in a breadth-first manner and emits events at specified + points. The events are handled by a visitor object that subclasses + :class:`~retworkx.visit.BFSVisitor` through the appropriate callback functions. + For example: + + .. jupyter-execute:: + + import retworkx + from retworkx.visit import BFSVisitor + + + class TreeEdgesRecorder(BFSVisitor): + + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append(edge) + + graph = retworkx.PyGraph() + graph.extend_from_edge_list([(1, 3), (0, 1), (2, 1), (0, 2)]) + vis = TreeEdgesRecorder() + retworkx.bfs_search(graph, [0], vis) + print('Tree edges:', vis.edges) \ No newline at end of file diff --git a/retworkx-core/src/lib.rs b/retworkx-core/src/lib.rs index 4d06732fe..cbee1fc55 100644 --- a/retworkx-core/src/lib.rs +++ b/retworkx-core/src/lib.rs @@ -70,11 +70,10 @@ pub type Result = core::result::Result; /// Module for centrality algorithms pub mod centrality; pub mod connectivity; -/// Module for depth first search edge methods -pub mod dfs_edges; /// Module for maximum weight matching algorithmss pub mod max_weight_matching; pub mod shortest_path; +pub mod traversal; // These modules define additional data structures pub mod dictmap; mod min_scored; diff --git a/retworkx-core/src/traversal/bfs_visit.rs b/retworkx-core/src/traversal/bfs_visit.rs new file mode 100644 index 000000000..4aca42d6e --- /dev/null +++ b/retworkx-core/src/traversal/bfs_visit.rs @@ -0,0 +1,282 @@ +// 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. + +use petgraph::visit::{ControlFlow, EdgeRef, IntoEdges, VisitMap, Visitable}; +use std::collections::VecDeque; + +/// Return if the expression is a break value, execute the provided statement +/// if it is a prune value. +/// https://github.com/petgraph/petgraph/blob/0.6.0/src/visit/dfsvisit.rs#L27 +macro_rules! try_control { + ($e:expr, $p:stmt) => { + try_control!($e, $p, ()); + }; + ($e:expr, $p:stmt, $q:stmt) => { + match $e { + x => { + if x.should_break() { + return x; + } else if x.should_prune() { + $p + } else { + $q + } + } + } + }; +} + +/// A breadth first search (BFS) visitor event. +#[derive(Copy, Clone, Debug)] +pub enum BfsEvent { + Discover(N), + /// An edge of the tree formed by the traversal. + TreeEdge(N, N, E), + /// An edge that does not belong to the tree. + NonTreeEdge(N, N, E), + /// For an edge *(u, v)*, if node *v* is currently in the queue + /// at the time of examination, then it is a gray-target edge. + GrayTargetEdge(N, N, E), + /// For an edge *(u, v)*, if node *v* has been removed from the queue + /// at the time of examination, then it is a black-target edge. + BlackTargetEdge(N, N, E), + /// All edges from a node have been reported. + Finish(N), +} + +/// An iterative breadth first search. +/// +/// Starting points are the nodes in the iterator `starts` (specify just one +/// start vertex *x* by using `Some(x)`). +/// +/// The traversal emits discovery and finish events for each reachable vertex, +/// and edge classification of each reachable edge. `visitor` is called for each +/// event, see [`BfsEvent`] for possible values. +/// +/// The return value should implement the trait [`ControlFlow`], and can be used to change +/// the control flow of the search. +/// +/// [`Control`](petgraph::visit::Control) Implements [`ControlFlow`] such that `Control::Continue` resumes the search. +/// `Control::Break` will stop the visit early, returning the contained value. +/// `Control::Prune` will stop traversing any additional edges from the current +/// node and proceed immediately to the `Finish` event. +/// +/// There are implementations of [`ControlFlow`] for `()`, and [`Result`] where +/// `C: ControlFlow`. The implementation for `()` will continue until finished. +/// For [`Result`], upon encountering an `E` it will break, otherwise acting the same as `C`. +/// +/// ***Panics** if you attempt to prune a node from its `Finish` event. +/// +/// The pseudo-code for the BFS algorithm is listed below, with the annotated +/// event points, for which the given visitor object will be called with the +/// appropriate method. +/// +/// ```norust +/// BFS(G, s) +/// for each vertex u in V +/// color[u] := WHITE +/// end for +/// color[s] := GRAY +/// EQUEUE(Q, s) discover vertex s +/// while (Q != Ø) +/// u := DEQUEUE(Q) +/// for each vertex v in Adj[u] (u,v) is a tree edge +/// if (color[v] = WHITE) +/// color[v] = GRAY +/// else (u,v) is a non - tree edge +/// if (color[v] = GRAY) (u,v) has a gray target +/// ... +/// else if (color[v] = BLACK) (u,v) has a black target +/// ... +/// end for +/// color[u] := BLACK finish vertex u +/// end while +/// ``` +/// +/// # Example returning [`Control`](petgraph::visit::Control). +/// +/// Find a path from vertex 0 to 5, and exit the visit as soon as we reach +/// the goal vertex. +/// +/// ``` +/// use retworkx_core::petgraph::prelude::*; +/// use retworkx_core::petgraph::graph::node_index as n; +/// use retworkx_core::petgraph::visit::Control; +/// +/// use retworkx_core::traversal::{BfsEvent, breadth_first_search}; +/// +/// let gr: Graph<(), ()> = Graph::from_edges(&[ +/// (0, 1), (0, 2), (0, 3), +/// (1, 3), +/// (2, 3), (2, 4), +/// (4, 0), (4, 5), +/// ]); +/// +/// // record each predecessor, mapping node → node +/// let mut predecessor = vec![NodeIndex::end(); gr.node_count()]; +/// let start = n(0); +/// let goal = n(5); +/// breadth_first_search(&gr, Some(start), |event| { +/// if let BfsEvent::TreeEdge(u, v, _) = event { +/// predecessor[v.index()] = u; +/// if v == goal { +/// return Control::Break(v); +/// } +/// } +/// Control::Continue +/// }); +/// +/// let mut next = goal; +/// let mut path = vec![next]; +/// while next != start { +/// let pred = predecessor[next.index()]; +/// path.push(pred); +/// next = pred; +/// } +/// path.reverse(); +/// assert_eq!(&path, &[n(0), n(2), n(4), n(5)]); +/// ``` +/// +/// # Example returning a `Result`. +/// ``` +/// use retworkx_core::petgraph::graph::node_index as n; +/// use retworkx_core::petgraph::prelude::*; +/// +/// use retworkx_core::traversal::{BfsEvent, breadth_first_search}; +/// +/// let gr: Graph<(), ()> = Graph::from_edges(&[(0, 1), (1, 2), (1, 1), (2, 1)]); +/// let start = n(0); +/// let mut non_tree_edges = 0; +/// +/// #[derive(Debug)] +/// struct NonTreeEdgeFound { +/// source: NodeIndex, +/// target: NodeIndex, +/// } +/// +/// // Stop the search, the first time a BackEdge is encountered. +/// let result = breadth_first_search(&gr, Some(start), |event| { +/// match event { +/// BfsEvent::NonTreeEdge(u, v, _) => { +/// non_tree_edges += 1; +/// // the implementation of ControlFlow for Result, +/// // treats this Err value as Continue::Break +/// Err(NonTreeEdgeFound {source: u, target: v}) +/// } +/// // In the cases where Ok(()) is returned, +/// // Result falls back to the implementation of Control on the value (). +/// // In the case of (), this is to always return Control::Continue. +/// // continuing the search. +/// _ => Ok(()), +/// } +/// }); +/// +/// assert_eq!(non_tree_edges, 1); +/// println!("number of non-tree edges encountered: {}", non_tree_edges); +/// println!("non-tree edge: ({:?})", result.unwrap_err()); +/// ``` +pub fn breadth_first_search( + graph: G, + starts: I, + mut visitor: F, +) -> C +where + G: IntoEdges + Visitable, + I: IntoIterator, + F: FnMut(BfsEvent) -> C, + C: ControlFlow, +{ + let discovered = &mut graph.visit_map(); + let finished = &mut graph.visit_map(); + + for start in starts { + // `bfs_visitor` returns a "signal" to either continue or exit early + // but it never "prunes", so we use `unreachable`. + try_control!( + bfs_visitor(graph, start, &mut visitor, discovered, finished), + unreachable!() + ); + } + C::continuing() +} + +fn bfs_visitor( + graph: G, + u: G::NodeId, + visitor: &mut F, + discovered: &mut G::Map, + finished: &mut G::Map, +) -> C +where + G: IntoEdges + Visitable, + F: FnMut(BfsEvent) -> C, + C: ControlFlow, +{ + if !discovered.visit(u) { + return C::continuing(); + } + + try_control!(visitor(BfsEvent::Discover(u)), {}, { + let mut stack: VecDeque = VecDeque::new(); + stack.push_front(u); + + while let Some(u) = stack.pop_front() { + for edge in graph.edges(u) { + let v = edge.target(); + if !discovered.is_visited(&v) { + try_control!( + visitor(BfsEvent::TreeEdge(u, v, edge.weight())), + continue + ); + discovered.visit(v); + try_control!(visitor(BfsEvent::Discover(v)), continue); + stack.push_back(v); + } else { + // non - tree edge. + try_control!( + visitor(BfsEvent::NonTreeEdge(u, v, edge.weight())), + continue + ); + + if !finished.is_visited(&v) { + try_control!( + visitor(BfsEvent::GrayTargetEdge( + u, + v, + edge.weight() + )), + continue + ); + } else { + try_control!( + visitor(BfsEvent::BlackTargetEdge( + u, + v, + edge.weight() + )), + continue + ); + } + } + } + + let first_finish = finished.visit(u); + debug_assert!(first_finish); + try_control!( + visitor(BfsEvent::Finish(u)), + panic!("Pruning on the `BfsEvent::Finish` is not supported!") + ); + } + }); + + C::continuing() +} diff --git a/retworkx-core/src/dfs_edges.rs b/retworkx-core/src/traversal/dfs_edges.rs similarity index 98% rename from retworkx-core/src/dfs_edges.rs rename to retworkx-core/src/traversal/dfs_edges.rs index 906cd8805..ec4b39b8c 100644 --- a/retworkx-core/src/dfs_edges.rs +++ b/retworkx-core/src/traversal/dfs_edges.rs @@ -32,7 +32,7 @@ use petgraph::visit::{ /// # Example /// ```rust /// use retworkx_core::petgraph; -/// use retworkx_core::dfs_edges::dfs_edges; +/// use retworkx_core::traversal::dfs_edges; /// /// let g = petgraph::graph::UnGraph::::from_edges(&[ /// (0, 1), (1, 2), (1, 3), (2, 4), (3, 4) diff --git a/retworkx-core/src/traversal/mod.rs b/retworkx-core/src/traversal/mod.rs new file mode 100644 index 000000000..2325fd9dc --- /dev/null +++ b/retworkx-core/src/traversal/mod.rs @@ -0,0 +1,19 @@ +// 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. + +//! Module for graph traversal algorithms. + +mod bfs_visit; +mod dfs_edges; + +pub use bfs_visit::{breadth_first_search, BfsEvent}; +pub use dfs_edges::dfs_edges; diff --git a/retworkx/__init__.py b/retworkx/__init__.py index 784be4bc2..47b479cc1 100644 --- a/retworkx/__init__.py +++ b/retworkx/__init__.py @@ -12,6 +12,9 @@ from .retworkx import * +# flake8: noqa +import retworkx.visit + sys.modules["retworkx.generators"] = generators @@ -1730,3 +1733,87 @@ def _graph_union( merge_edges=False, ): return graph_union(first, second, merge_nodes=False, merge_edges=False) + + +@functools.singledispatch +def bfs_search(graph, source, visitor): + """Breadth-first traversal of a directed/undirected graph. + + The pseudo-code for the BFS algorithm is listed below, with the annotated + event points, for which the given visitor object will be called with the + appropriate method. + + :: + + BFS(G, s) + for each vertex u in V + color[u] := WHITE + end for + color[s] := GRAY + EQUEUE(Q, s) discover vertex s + while (Q != Ø) + u := DEQUEUE(Q) + for each vertex v in Adj[u] (u,v) is a tree edge + if (color[v] = WHITE) + color[v] = GRAY + else (u,v) is a non - tree edge + if (color[v] = GRAY) (u,v) has a gray target + ... + else if (color[v] = BLACK) (u,v) has a black target + ... + end for + color[u] := BLACK finish vertex u + end while + + If an exception is raised inside the callback function, the graph traversal + will be stopped immediately. You can exploit this to exit early by raising a + :class:`~retworkx.visit.StopSearch` exception, in which case the search function + will return but without raising back the exception. You can also prune part of + the search tree by raising :class:`~retworkx.visit.PruneSearch`. + + In the following example we keep track of the tree edges: + + .. jupyter-execute:: + + import retworkx + from retworkx.visit import BFSVisitor + + + class TreeEdgesRecorder(BFSVisitor): + + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append(edge) + + graph = retworkx.PyDiGraph() + graph.extend_from_edge_list([(1, 3), (0, 1), (2, 1), (0, 2)]) + vis = TreeEdgesRecorder() + retworkx.bfs_search(graph, [0], vis) + print('Tree edges:', vis.edges) + + .. note:: + + Graph can **not** be mutated while traversing. + + :param graph: The graph to be used. This can be a :class:`~retworkx.PyGraph` + or a :class:`~retworkx.PyDiGraph` + :param List[int] source: An optional list of node indices to use as the starting + nodes for the breadth-first search. If this is not specified then a source + will be chosen arbitrarly and repeated until all components of the + graph are searched. + :param visitor: A visitor object that is invoked at the event points inside the + algorithm. This should be a subclass of :class:`~retworkx.visit.BFSVisitor`. + """ + raise TypeError("Invalid Input Type %s for graph" % type(graph)) + + +@bfs_search.register(PyDiGraph) +def _digraph_bfs_search(graph, source, visitor): + return digraph_bfs_search(graph, source, visitor) + + +@bfs_search.register(PyGraph) +def _graph_bfs_search(graph, source, visitor): + return graph_bfs_search(graph, source, visitor) diff --git a/retworkx/visit.py b/retworkx/visit.py new file mode 100644 index 000000000..8c59d727a --- /dev/null +++ b/retworkx/visit.py @@ -0,0 +1,71 @@ +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + + +class StopSearch(Exception): + """Stop graph traversal""" + + pass + + +class PruneSearch(Exception): + """Prune part of the search tree while traversing a graph.""" + + pass + + +class BFSVisitor: + """A visitor object that is invoked at the event-points inside the + :func:`~retworkx.bfs_search` algorithm. By default, it performs no + action, and should be used as a base class in order to be useful. + """ + + def discover_vertex(self, v): + """ + This is invoked when a vertex is encountered for the first time. + """ + return + + def finish_vertex(self, v): + """ + This is invoked on vertex `v` after all of its out edges have been + added to the search tree and all of the adjacent vertices have been + discovered, but before the out-edges of the adjacent vertices have + been examined. + """ + return + + def tree_edge(self, e): + """ + This is invoked on each edge as it becomes a member of the edges + that form the search tree. + """ + return + + def non_tree_edge(self, e): + """ + This is invoked on back or cross edges for directed graphs and cross edges + for undirected graphs. + """ + return + + def gray_target_edge(self, e): + """ + This is invoked on the subset of non-tree edges whose target vertex is + colored gray at the time of examination. + The color gray indicates that the vertex is currently in the queue. + """ + return + + def black_target_edge(self, e): + """ + This is invoked on the subset of non-tree edges whose target vertex is + colored black at the time of examination. + The color black indicates that the vertex has been removed from the queue. + """ + return diff --git a/src/lib.rs b/src/lib.rs index cad84e94c..1334ef628 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -50,6 +50,7 @@ use num_complex::Complex64; use pyo3::create_exception; use pyo3::exceptions::PyException; +use pyo3::import_exception; use pyo3::prelude::*; use pyo3::wrap_pyfunction; use pyo3::wrap_pymodule; @@ -198,6 +199,10 @@ create_exception!(retworkx, NoSuitableNeighbors, PyException); create_exception!(retworkx, NullGraph, PyException); // No path was found between the specified nodes. create_exception!(retworkx, NoPathFound, PyException); +// Prune part of the search tree while traversing a graph. +import_exception!(retworkx.visit, PruneSearch); +// Stop graph traversal. +import_exception!(retworkx.visit, StopSearch); #[pymodule] fn retworkx(py: Python<'_>, m: &PyModule) -> PyResult<()> { @@ -210,6 +215,8 @@ fn retworkx(py: Python<'_>, m: &PyModule) -> PyResult<()> { m.add("NoPathFound", py.get_type::())?; m.add("NullGraph", py.get_type::())?; m.add_wrapped(wrap_pyfunction!(bfs_successors))?; + m.add_wrapped(wrap_pyfunction!(graph_bfs_search))?; + m.add_wrapped(wrap_pyfunction!(digraph_bfs_search))?; m.add_wrapped(wrap_pyfunction!(dag_longest_path))?; m.add_wrapped(wrap_pyfunction!(dag_longest_path_length))?; m.add_wrapped(wrap_pyfunction!(dag_weighted_longest_path))?; diff --git a/src/traversal/bfs_visit.rs b/src/traversal/bfs_visit.rs new file mode 100644 index 000000000..446931cb4 --- /dev/null +++ b/src/traversal/bfs_visit.rs @@ -0,0 +1,69 @@ +// 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. + +use pyo3::prelude::*; + +use petgraph::stable_graph::NodeIndex; +use petgraph::visit::Control; + +use crate::{PruneSearch, StopSearch}; +use retworkx_core::traversal::BfsEvent; + +#[derive(FromPyObject)] +pub struct PyBfsVisitor { + discover_vertex: PyObject, + finish_vertex: PyObject, + tree_edge: PyObject, + non_tree_edge: PyObject, + gray_target_edge: PyObject, + black_target_edge: PyObject, +} + +pub fn bfs_handler( + py: Python, + vis: &PyBfsVisitor, + event: BfsEvent, +) -> PyResult> { + let res = match event { + BfsEvent::Discover(u) => vis.discover_vertex.call1(py, (u.index(),)), + BfsEvent::TreeEdge(u, v, weight) => { + let edge = (u.index(), v.index(), weight); + vis.tree_edge.call1(py, (edge,)) + } + BfsEvent::NonTreeEdge(u, v, weight) => { + let edge = (u.index(), v.index(), weight); + vis.non_tree_edge.call1(py, (edge,)) + } + BfsEvent::GrayTargetEdge(u, v, weight) => { + let edge = (u.index(), v.index(), weight); + vis.gray_target_edge.call1(py, (edge,)) + } + BfsEvent::BlackTargetEdge(u, v, weight) => { + let edge = (u.index(), v.index(), weight); + vis.black_target_edge.call1(py, (edge,)) + } + BfsEvent::Finish(u) => vis.finish_vertex.call1(py, (u.index(),)), + }; + + match res { + Err(e) => { + if e.is_instance::(py) { + Ok(Control::Prune) + } else if e.is_instance::(py) { + Ok(Control::Break(())) + } else { + Err(e) + } + } + Ok(_) => Ok(Control::Continue), + } +} diff --git a/src/traversal/mod.rs b/src/traversal/mod.rs index 53d4936ec..71348fcdd 100644 --- a/src/traversal/mod.rs +++ b/src/traversal/mod.rs @@ -10,7 +10,10 @@ // License for the specific language governing permissions and limitations // under the License. -use retworkx_core::dfs_edges; +mod bfs_visit; + +use bfs_visit::{bfs_handler, PyBfsVisitor}; +use retworkx_core::traversal::{breadth_first_search, dfs_edges}; use super::{digraph, graph, iterators}; @@ -44,7 +47,7 @@ fn digraph_dfs_edges( source: Option, ) -> EdgeList { EdgeList { - edges: dfs_edges::dfs_edges(&graph.graph, source.map(NodeIndex::new)), + edges: dfs_edges(&graph.graph, source.map(NodeIndex::new)), } } @@ -64,7 +67,7 @@ fn digraph_dfs_edges( #[pyo3(text_signature = "(graph, /, source=None)")] fn graph_dfs_edges(graph: &graph::PyGraph, source: Option) -> EdgeList { EdgeList { - edges: dfs_edges::dfs_edges(&graph.graph, source.map(NodeIndex::new)), + edges: dfs_edges(&graph.graph, source.map(NodeIndex::new)), } } @@ -165,3 +168,175 @@ fn descendants(graph: &digraph::PyDiGraph, node: usize) -> HashSet { out_set.remove(&node); out_set } + +/// Breadth-first traversal of a directed graph. +/// +/// The pseudo-code for the BFS algorithm is listed below, with the annotated +/// event points, for which the given visitor object will be called with the +/// appropriate method. +/// +/// :: +/// +/// BFS(G, s) +/// for each vertex u in V +/// color[u] := WHITE +/// end for +/// color[s] := GRAY +/// EQUEUE(Q, s) discover vertex s +/// while (Q != Ø) +/// u := DEQUEUE(Q) +/// for each vertex v in Adj[u] (u,v) is a tree edge +/// if (color[v] = WHITE) +/// color[v] = GRAY +/// else (u,v) is a non - tree edge +/// if (color[v] = GRAY) (u,v) has a gray target +/// ... +/// else if (color[v] = BLACK) (u,v) has a black target +/// ... +/// end for +/// color[u] := BLACK finish vertex u +/// end while +/// +/// If an exception is raised inside the callback function, the graph traversal +/// will be stopped immediately. You can exploit this to exit early by raising a +/// :class:`~retworkx.visit.StopSearch` exception, in which case the search function +/// will return but without raising back the exception. You can also prune part of the +/// search tree by raising :class:`~retworkx.visit.PruneSearch`. +/// +/// In the following example we keep track of the tree edges: +/// +/// .. jupyter-execute:: +/// +/// import retworkx +/// from retworkx.visit import BFSVisitor +/// +/// class TreeEdgesRecorder(BFSVisitor): +/// +/// def __init__(self): +/// self.edges = [] +/// +/// def tree_edge(self, edge): +/// self.edges.append(edge) +/// +/// graph = retworkx.PyDiGraph() +/// graph.extend_from_edge_list([(1, 3), (0, 1), (2, 1), (0, 2)]) +/// vis = TreeEdgesRecorder() +/// retworkx.bfs_search(graph, [0], vis) +/// print('Tree edges:', vis.edges) +/// +/// .. note:: +/// +/// Graph can **not** be mutated while traversing. +/// +/// :param PyDiGraph graph: The graph to be used. +/// :param List[int] source: An optional list of node indices to use as the starting nodes +/// for the breadth-first search. If this is not specified then a source +/// will be chosen arbitrarly and repeated until all components of the +/// graph are searched. +/// :param visitor: A visitor object that is invoked at the event points inside the +/// algorithm. This should be a subclass of :class:`~retworkx.visit.BFSVisitor`. +#[pyfunction] +#[pyo3(text_signature = "(graph, source, visitor)")] +pub fn digraph_bfs_search( + py: Python, + graph: &digraph::PyDiGraph, + source: Option>, + visitor: PyBfsVisitor, +) -> PyResult<()> { + let starts: Vec<_> = match source { + Some(nx) => nx.into_iter().map(NodeIndex::new).collect(), + None => graph.graph.node_indices().collect(), + }; + + breadth_first_search(&graph.graph, starts, |event| { + bfs_handler(py, &visitor, event) + })?; + + Ok(()) +} + +/// Breadth-first traversal of an undirected graph. +/// +/// The pseudo-code for the BFS algorithm is listed below, with the annotated +/// event points, for which the given visitor object will be called with the +/// appropriate method. +/// +/// :: +/// +/// BFS(G, s) +/// for each vertex u in V +/// color[u] := WHITE +/// end for +/// color[s] := GRAY +/// EQUEUE(Q, s) discover vertex s +/// while (Q != Ø) +/// u := DEQUEUE(Q) +/// for each vertex v in Adj[u] (u,v) is a tree edge +/// if (color[v] = WHITE) +/// color[v] = GRAY +/// else (u,v) is a non - tree edge +/// if (color[v] = GRAY) (u,v) has a gray target +/// ... +/// else if (color[v] = BLACK) (u,v) has a black target +/// ... +/// end for +/// color[u] := BLACK finish vertex u +/// end while +/// +/// If an exception is raised inside the callback function, the graph traversal +/// will be stopped immediately. You can exploit this to exit early by raising a +/// :class:`~retworkx.visit.StopSearch` exception, in which case the search function +/// will return but without raising back the exception. You can also prune part of the +/// search tree by raising :class:`~retworkx.visit.PruneSearch`. +/// +/// In the following example we keep track of the tree edges: +/// +/// .. jupyter-execute:: +/// +/// import retworkx +/// from retworkx.visit import BFSVisitor +/// +/// class TreeEdgesRecorder(BFSVisitor): +/// +/// def __init__(self): +/// self.edges = [] +/// +/// def tree_edge(self, edge): +/// self.edges.append(edge) +/// +/// graph = retworkx.PyGraph() +/// graph.extend_from_edge_list([(1, 3), (0, 1), (2, 1), (0, 2)]) +/// vis = TreeEdgesRecorder() +/// retworkx.bfs_search(graph, [0], vis) +/// print('Tree edges:', vis.edges) +/// +/// .. note:: +/// +/// Graph can **not** be mutated while traversing. +/// +/// :param PyGraph graph: The graph to be used. +/// :param List[int] source: An optional list of node indices to use as the starting nodes +/// for the breadth-first search. If this is not specified then a source +/// will be chosen arbitrarly and repeated until all components of the +/// graph are searched. +/// :param visitor: A visitor object that is invoked at the event points inside the +/// algorithm. This should be a subclass of :class:`~retworkx.visit.BFSVisitor`. +#[pyfunction] +#[pyo3(text_signature = "(graph, source, visitor)")] +pub fn graph_bfs_search( + py: Python, + graph: &graph::PyGraph, + source: Option>, + visitor: PyBfsVisitor, +) -> PyResult<()> { + let starts: Vec<_> = match source { + Some(nx) => nx.into_iter().map(NodeIndex::new).collect(), + None => graph.graph.node_indices().collect(), + }; + + breadth_first_search(&graph.graph, starts, |event| { + bfs_handler(py, &visitor, event) + })?; + + Ok(()) +} diff --git a/tests/digraph/test_bfs_search.py b/tests/digraph/test_bfs_search.py new file mode 100644 index 000000000..1ada863a5 --- /dev/null +++ b/tests/digraph/test_bfs_search.py @@ -0,0 +1,164 @@ +# 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 retworkx + + +class TestBfsSearch(unittest.TestCase): + def setUp(self): + self.graph = retworkx.PyDiGraph() + self.graph.extend_from_edge_list( + [ + (0, 1), + (0, 2), + (1, 3), + (2, 1), + (2, 5), + (2, 6), + (5, 3), + (4, 7), + ] + ) + + def test_digraph_bfs_tree_edges(self): + class TreeEdgesRecorder(retworkx.visit.BFSVisitor): + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append((edge[0], edge[1])) + + vis = TreeEdgesRecorder() + retworkx.digraph_bfs_search(self.graph, [0], vis) + self.assertEqual(vis.edges, [(0, 2), (0, 1), (2, 6), (2, 5), (1, 3)]) + + def test_digraph_bfs_tree_edges_no_starting_point(self): + class TreeEdgesRecorder(retworkx.visit.BFSVisitor): + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append((edge[0], edge[1])) + + vis = TreeEdgesRecorder() + retworkx.digraph_bfs_search(self.graph, None, vis) + self.assertEqual( + vis.edges, [(0, 2), (0, 1), (2, 6), (2, 5), (1, 3), (4, 7)] + ) + + def test_digraph_bfs_tree_edges_restricted(self): + class TreeEdgesRecorderRestricted(retworkx.visit.BFSVisitor): + + prohibited = [(0, 2), (1, 2)] + + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + edge = (edge[0], edge[1]) + if edge in self.prohibited: + raise retworkx.visit.PruneSearch + self.edges.append(edge) + + vis = TreeEdgesRecorderRestricted() + retworkx.digraph_bfs_search(self.graph, [0], vis) + self.assertEqual(vis.edges, [(0, 1), (1, 3)]) + + def test_digraph_bfs_goal_search_with_stop_search_exception(self): + class GoalSearch(retworkx.visit.BFSVisitor): + + goal = 3 + + def __init__(self): + self.parents = {} + + def tree_edge(self, edge): + u, v, _ = edge + self.parents[v] = u + + if v == self.goal: + raise retworkx.visit.StopSearch + + def reconstruct_path(self): + v = self.goal + path = [v] + while v in self.parents: + v = self.parents[v] + path.append(v) + + path.reverse() + return path + + vis = GoalSearch() + retworkx.digraph_bfs_search(self.graph, [0], vis) + self.assertEqual(vis.reconstruct_path(), [0, 1, 3]) + + def test_digraph_bfs_goal_search_with_custom_exception(self): + class StopIfGoalFound(Exception): + pass + + class GoalSearch(retworkx.visit.BFSVisitor): + + goal = 3 + + def __init__(self): + self.parents = {} + + def tree_edge(self, edge): + u, v, _ = edge + self.parents[v] = u + + if v == self.goal: + raise StopIfGoalFound + + def reconstruct_path(self): + v = self.goal + path = [v] + while v in self.parents: + v = self.parents[v] + path.append(v) + + path.reverse() + return path + + vis = GoalSearch() + try: + retworkx.digraph_bfs_search(self.graph, [0], vis) + except StopIfGoalFound: + pass + self.assertEqual(vis.reconstruct_path(), [0, 1, 3]) + + def test_graph_prune_non_tree_edge(self): + class PruneNonTreeEdge(retworkx.visit.BFSVisitor): + def non_tree_edge(self, _): + raise retworkx.visit.PruneSearch + + vis = PruneNonTreeEdge() + retworkx.digraph_bfs_search(self.graph, [0], vis) + + def test_graph_prune_black_target_edge(self): + class PruneBlackTargetEdge(retworkx.visit.BFSVisitor): + def black_target_edge(self, _): + raise retworkx.visit.PruneSearch + + vis = PruneBlackTargetEdge() + retworkx.digraph_bfs_search(self.graph, [0], vis) + + def test_graph_prune_gray_target_edge(self): + class PruneGrayTargetEdge(retworkx.visit.BFSVisitor): + def gray_target_edge(self, _): + raise retworkx.visit.PruneSearch + + vis = PruneGrayTargetEdge() + retworkx.digraph_bfs_search(self.graph, [0], vis) diff --git a/tests/graph/test_bfs_search.py b/tests/graph/test_bfs_search.py new file mode 100644 index 000000000..db6c7fc49 --- /dev/null +++ b/tests/graph/test_bfs_search.py @@ -0,0 +1,164 @@ +# 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 retworkx + + +class TestBfsSearch(unittest.TestCase): + def setUp(self): + self.graph = retworkx.PyGraph() + self.graph.extend_from_edge_list( + [ + (0, 1), + (0, 2), + (1, 3), + (2, 1), + (2, 5), + (2, 6), + (5, 3), + (4, 7), + ] + ) + + def test_graph_bfs_tree_edges(self): + class TreeEdgesRecorder(retworkx.visit.BFSVisitor): + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append((edge[0], edge[1])) + + vis = TreeEdgesRecorder() + retworkx.graph_bfs_search(self.graph, [0], vis) + self.assertEqual(vis.edges, [(0, 2), (0, 1), (2, 6), (2, 5), (1, 3)]) + + def test_graph_bfs_tree_edges_no_starting_point(self): + class TreeEdgesRecorder(retworkx.visit.BFSVisitor): + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append((edge[0], edge[1])) + + vis = TreeEdgesRecorder() + retworkx.graph_bfs_search(self.graph, None, vis) + self.assertEqual( + vis.edges, [(0, 2), (0, 1), (2, 6), (2, 5), (1, 3), (4, 7)] + ) + + def test_graph_bfs_tree_edges_restricted(self): + class TreeEdgesRecorderRestricted(retworkx.visit.BFSVisitor): + + prohibited = [(0, 2), (1, 2)] + + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + edge = (edge[0], edge[1]) + if edge in self.prohibited: + raise retworkx.visit.PruneSearch + self.edges.append(edge) + + vis = TreeEdgesRecorderRestricted() + retworkx.graph_bfs_search(self.graph, [0], vis) + self.assertEqual(vis.edges, [(0, 1), (1, 3), (3, 5), (5, 2), (2, 6)]) + + def test_graph_bfs_goal_search_with_stop_search_exception(self): + class GoalSearch(retworkx.visit.BFSVisitor): + + goal = 3 + + def __init__(self): + self.parents = {} + + def tree_edge(self, edge): + u, v, _ = edge + self.parents[v] = u + + if v == self.goal: + raise retworkx.visit.StopSearch + + def reconstruct_path(self): + v = self.goal + path = [v] + while v in self.parents: + v = self.parents[v] + path.append(v) + + path.reverse() + return path + + vis = GoalSearch() + retworkx.graph_bfs_search(self.graph, [0], vis) + self.assertEqual(vis.reconstruct_path(), [0, 1, 3]) + + def test_graph_bfs_goal_search_with_custom_exception(self): + class StopIfGoalFound(Exception): + pass + + class GoalSearch(retworkx.visit.BFSVisitor): + + goal = 3 + + def __init__(self): + self.parents = {} + + def tree_edge(self, edge): + u, v, _ = edge + self.parents[v] = u + + if v == self.goal: + raise StopIfGoalFound + + def reconstruct_path(self): + v = self.goal + path = [v] + while v in self.parents: + v = self.parents[v] + path.append(v) + + path.reverse() + return path + + vis = GoalSearch() + try: + retworkx.graph_bfs_search(self.graph, [0], vis) + except StopIfGoalFound: + pass + self.assertEqual(vis.reconstruct_path(), [0, 1, 3]) + + def test_graph_prune_non_tree_edge(self): + class PruneNonTreeEdge(retworkx.visit.BFSVisitor): + def non_tree_edge(self, _): + raise retworkx.visit.PruneSearch + + vis = PruneNonTreeEdge() + retworkx.graph_bfs_search(self.graph, [0], vis) + + def test_graph_prune_black_target_edge(self): + class PruneBlackTargetEdge(retworkx.visit.BFSVisitor): + def black_target_edge(self, _): + raise retworkx.visit.PruneSearch + + vis = PruneBlackTargetEdge() + retworkx.graph_bfs_search(self.graph, [0], vis) + + def test_graph_prune_gray_target_edge(self): + class PruneGrayTargetEdge(retworkx.visit.BFSVisitor): + def gray_target_edge(self, _): + raise retworkx.visit.PruneSearch + + vis = PruneGrayTargetEdge() + retworkx.graph_bfs_search(self.graph, [0], vis)