diff --git a/docs/source/api.rst b/docs/source/api.rst index 9b219aac2..fd7740933 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -56,6 +56,7 @@ Traversal :toctree: stubs retworkx.dfs_edges + retworkx.dfs_search retworkx.bfs_successors retworkx.bfs_search retworkx.topological_sort @@ -64,6 +65,7 @@ Traversal retworkx.ancestors retworkx.collect_runs retworkx.collect_bicolor_runs + retworkx.visit.DFSVisitor retworkx.visit.BFSVisitor .. _dag-algorithms: @@ -250,6 +252,7 @@ the functions from the explicitly typed based on the data type. retworkx.digraph_all_pairs_dijkstra_path_lengths retworkx.digraph_k_shortest_path_lengths retworkx.digraph_dfs_edges + retworkx.digraph_dfs_search retworkx.digraph_find_cycle retworkx.digraph_transitivity retworkx.digraph_core_number @@ -295,6 +298,7 @@ typed API based on the data type. retworkx.graph_k_shortest_path_lengths retworkx.graph_all_pairs_dijkstra_path_lengths retworkx.graph_dfs_edges + retworkx.graph_dfs_search retworkx.graph_transitivity retworkx.graph_core_number retworkx.graph_complement diff --git a/releasenotes/notes/dfs-search-6083680bf62356b0.yaml b/releasenotes/notes/dfs-search-6083680bf62356b0.yaml new file mode 100644 index 000000000..c2d74e721 --- /dev/null +++ b/releasenotes/notes/dfs-search-6083680bf62356b0.yaml @@ -0,0 +1,28 @@ +--- +features: + - | + Added a new :func:`~retworkx.dfs_search` (and it's per type variants + :func:`~retworkx.graph_dfs_search` and :func:`~retworkx.digraph_dfs_search`) + that traverses the graph in a depth-first manner and emits events at specified + points. The events are handled by a visitor object that subclasses + :class:`~retworkx.visit.DFSVisitor` through the appropriate callback functions. + For example: + + .. jupyter-execute:: + + import retworkx + from retworkx.visit import DFSVisitor + + class TreeEdgesRecorder(DFSVisitor): + + 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.dfs_search(graph, [0], vis) + print('Tree edges:', vis.edges) diff --git a/retworkx-core/src/traversal/bfs_visit.rs b/retworkx-core/src/traversal/bfs_visit.rs index 4aca42d6e..97ab2ab9c 100644 --- a/retworkx-core/src/traversal/bfs_visit.rs +++ b/retworkx-core/src/traversal/bfs_visit.rs @@ -13,27 +13,7 @@ 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 - } - } - } - }; -} +use super::try_control; /// A breadth first search (BFS) visitor event. #[derive(Copy, Clone, Debug)] diff --git a/retworkx-core/src/traversal/dfs_visit.rs b/retworkx-core/src/traversal/dfs_visit.rs new file mode 100644 index 000000000..be7e0c575 --- /dev/null +++ b/retworkx-core/src/traversal/dfs_visit.rs @@ -0,0 +1,260 @@ +// 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 is an iterative implementation of the upstream petgraph +// ``depth_first_search`` function. +// https://github.com/petgraph/petgraph/blob/0.6.0/src/visit/dfsvisit.rs + +use petgraph::visit::{ + ControlFlow, EdgeRef, IntoEdges, Time, VisitMap, Visitable, +}; + +use super::try_control; + +/// A depth first search (DFS) visitor event. +/// +/// It's similar to upstream petgraph +/// [`DfsEvent`](https://docs.rs/petgraph/0.6.0/petgraph/visit/enum.DfsEvent.html) +/// event. +#[derive(Copy, Clone, Debug)] +pub enum DfsEvent { + Discover(N, Time), + /// An edge of the tree formed by the traversal. + TreeEdge(N, N, E), + /// An edge to an already visited node. + BackEdge(N, N, E), + /// A cross or forward edge. + /// + /// For an edge *(u, v)*, if the discover time of *v* is greater than *u*, + /// then it is a forward edge, else a cross edge. + CrossForwardEdge(N, N, E), + /// All edges from a node have been reported. + Finish(N, Time), +} + +/// An iterative depth first search. +/// +/// It is an iterative implementation of the upstream petgraph +/// [`depth_first_search`](https://docs.rs/petgraph/0.6.0/petgraph/visit/fn.depth_first_search.html) function. +/// +/// 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 `petgraph::DfsEvent` for possible values. +/// +/// The return value should implement the trait `ControlFlow`, and can be used to change +/// the control flow of the search. +/// +/// `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. +/// +/// # Example returning `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::{DfsEvent, depth_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); +/// depth_first_search(&gr, Some(start), |event| { +/// if let DfsEvent::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::petgraph::visit::Time; +/// +/// use retworkx_core::traversal::{DfsEvent, depth_first_search}; +/// +/// let gr: Graph<(), ()> = Graph::from_edges(&[(0, 1), (1, 2), (1, 1), (2, 1)]); +/// let start = n(0); +/// let mut back_edges = 0; +/// let mut discover_time = 0; +/// +/// #[derive(Debug)] +/// struct BackEdgeFound { +/// source: NodeIndex, +/// target: NodeIndex, +/// } +/// +/// // Stop the search, the first time a BackEdge is encountered. +/// let result = depth_first_search(&gr, Some(start), |event| { +/// match event { +/// // 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. +/// DfsEvent::Discover(_, Time(t)) => { +/// discover_time = t; +/// Ok(()) +/// } +/// DfsEvent::BackEdge(u, v, _) => { +/// back_edges += 1; +/// // the implementation of ControlFlow for Result, +/// // treats this Err value as Continue::Break +/// Err(BackEdgeFound {source: u, target: v}) +/// } +/// _ => Ok(()), +/// } +/// }); +/// +/// // Even though the graph has more than one cycle, +/// // The number of back_edges visited by the search should always be 1. +/// assert_eq!(back_edges, 1); +/// println!("discover time:{:?}", discover_time); +/// println!("number of backedges encountered: {}", back_edges); +/// println!("back edge: ({:?})", result.unwrap_err()); +/// ``` +pub fn depth_first_search(graph: G, starts: I, mut visitor: F) -> C +where + G: IntoEdges + Visitable, + I: IntoIterator, + F: FnMut(DfsEvent) -> C, + C: ControlFlow, +{ + let time = &mut Time(0); + let discovered = &mut graph.visit_map(); + let finished = &mut graph.visit_map(); + + for start in starts { + try_control!( + dfs_visitor(graph, start, &mut visitor, discovered, finished, time), + unreachable!() + ); + } + C::continuing() +} + +fn dfs_visitor( + graph: G, + u: G::NodeId, + visitor: &mut F, + discovered: &mut G::Map, + finished: &mut G::Map, + time: &mut Time, +) -> C +where + G: IntoEdges + Visitable, + F: FnMut(DfsEvent) -> C, + C: ControlFlow, +{ + if !discovered.visit(u) { + return C::continuing(); + } + + try_control!(visitor(DfsEvent::Discover(u, time_post_inc(time))), {}, { + let mut stack: Vec<(G::NodeId, ::Edges)> = Vec::new(); + stack.push((u, graph.edges(u))); + + while let Some(elem) = stack.last_mut() { + let u = elem.0; + let adjacent_edges = &mut elem.1; + let mut next = None; + + for edge in adjacent_edges { + let v = edge.target(); + if !discovered.is_visited(&v) { + try_control!( + visitor(DfsEvent::TreeEdge(u, v, edge.weight())), + continue + ); + discovered.visit(v); + try_control!( + visitor(DfsEvent::Discover(v, time_post_inc(time))), + continue + ); + next = Some(v); + break; + } else if !finished.is_visited(&v) { + try_control!( + visitor(DfsEvent::BackEdge(u, v, edge.weight())), + continue + ); + } else { + try_control!( + visitor(DfsEvent::CrossForwardEdge( + u, + v, + edge.weight() + )), + continue + ); + } + } + + match next { + Some(v) => stack.push((v, graph.edges(v))), + None => { + let first_finish = finished.visit(u); + debug_assert!(first_finish); + try_control!( + visitor(DfsEvent::Finish(u, time_post_inc(time))), + panic!("Pruning on the `DfsEvent::Finish` is not supported!") + ); + stack.pop(); + } + }; + } + }); + + C::continuing() +} + +fn time_post_inc(x: &mut Time) -> Time { + let v = *x; + x.0 += 1; + v +} diff --git a/retworkx-core/src/traversal/mod.rs b/retworkx-core/src/traversal/mod.rs index 2325fd9dc..2a06c3007 100644 --- a/retworkx-core/src/traversal/mod.rs +++ b/retworkx-core/src/traversal/mod.rs @@ -14,6 +14,32 @@ mod bfs_visit; mod dfs_edges; +mod dfs_visit; pub use bfs_visit::{breadth_first_search, BfsEvent}; pub use dfs_edges::dfs_edges; +pub use dfs_visit::{depth_first_search, DfsEvent}; + +/// 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 + } + } + } + }; +} + +use try_control; diff --git a/retworkx/__init__.py b/retworkx/__init__.py index 641bf109c..79b62eb86 100644 --- a/retworkx/__init__.py +++ b/retworkx/__init__.py @@ -1892,3 +1892,83 @@ def _digraph_bfs_search(graph, source, visitor): @bfs_search.register(PyGraph) def _graph_bfs_search(graph, source, visitor): return graph_bfs_search(graph, source, visitor) + + +@functools.singledispatch +def dfs_search(graph, source, visitor): + """Depth-first traversal of a directed/undirected graph. + + The pseudo-code for the DFS algorithm is listed below, with the annotated + event points, for which the given visitor object will be called with the + appropriate method. + + :: + + DFS(G) + for each vertex u in V + color[u] := WHITE initialize vertex u + end for + time := 0 + call DFS-VISIT(G, source) start vertex s + + DFS-VISIT(G, u) + color[u] := GRAY discover vertex u + for each v in Adj[u] examine edge (u,v) + if (color[v] = WHITE) (u,v) is a tree edge + all DFS-VISIT(G, v) + else if (color[v] = GRAY) (u,v) is a back edge + ... + else if (color[v] = BLACK) (u,v) is a cross or forward edge + ... + end for + color[u] := BLACK finish vertex u + + 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. 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 DFSVisitor + + class TreeEdgesRecorder(DFSVisitor): + + 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.dfs_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 depth-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.DFSVisitor`. + """ + raise TypeError("Invalid Input Type %s for graph" % type(graph)) + + +@dfs_search.register(PyDiGraph) +def _digraph_dfs_search(graph, source, visitor): + return digraph_dfs_search(graph, source, visitor) + + +@dfs_search.register(PyGraph) +def _graph_dfs_search(graph, source, visitor): + return graph_dfs_search(graph, source, visitor) diff --git a/retworkx/visit.py b/retworkx/visit.py index 8c59d727a..d0f21fb68 100644 --- a/retworkx/visit.py +++ b/retworkx/visit.py @@ -69,3 +69,54 @@ def black_target_edge(self, e): The color black indicates that the vertex has been removed from the queue. """ return + + +class DFSVisitor: + """A visitor object that is invoked at the event-points inside the + :func:`~retworkx.dfs_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, t): + """ + This is invoked when a vertex is encountered for the first time. + Together we report the discover time of vertex `v`. + """ + return + + def finish_vertex(self, v, t): + """ + This is invoked on vertex `v` after `finish_vertex` has been called for all + the vertices in the DFS-tree rooted at vertex `v`. If vertex `v` is a leaf in + the DFS-tree, then the `finish_vertex` function is called on `v` after all + the out-edges of `v` have been examined. Together we report the finish time + of vertex `v`. + """ + 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 back_edge(self, e): + """ + This is invoked on the back edges in the graph. + For an undirected graph there is some ambiguity between tree edges + and back edges since the edge :math:`(u, v)` and :math:`(v, u)` are the + same edge, but both the `tree_edge()` and `back_edge()` functions will be + invoked. One way to resolve this ambiguity is to record the tree edges, + and then disregard the back-edges that are already marked as tree edges. + An easy way to record tree edges is to record predecessors at the + `tree_edge` event point. + """ + return + + def forward_or_cross_edge(self, e): + """ + This is invoked on forward or cross edges in the graph. + In an undirected graph this method is never called. + """ + return diff --git a/src/lib.rs b/src/lib.rs index 39bb36725..30c0ac3fe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -405,6 +405,8 @@ fn retworkx(py: Python<'_>, m: &PyModule) -> PyResult<()> { ))?; m.add_wrapped(wrap_pyfunction!(metric_closure))?; m.add_wrapped(wrap_pyfunction!(steiner_tree))?; + m.add_wrapped(wrap_pyfunction!(digraph_dfs_search))?; + m.add_wrapped(wrap_pyfunction!(graph_dfs_search))?; m.add_wrapped(wrap_pyfunction!(chain_decomposition))?; m.add_class::()?; m.add_class::()?; diff --git a/src/traversal/dfs_visit.rs b/src/traversal/dfs_visit.rs new file mode 100644 index 000000000..8d53a88cc --- /dev/null +++ b/src/traversal/dfs_visit.rs @@ -0,0 +1,66 @@ +// 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, Time}; + +use crate::PruneSearch; +use retworkx_core::traversal::DfsEvent; + +#[derive(FromPyObject)] +pub struct PyDfsVisitor { + discover_vertex: PyObject, + finish_vertex: PyObject, + tree_edge: PyObject, + back_edge: PyObject, + forward_or_cross_edge: PyObject, +} + +pub fn dfs_handler( + py: Python, + vis: &PyDfsVisitor, + event: DfsEvent, +) -> PyResult> { + let res = match event { + DfsEvent::Discover(u, Time(t)) => { + vis.discover_vertex.call1(py, (u.index(), t)) + } + DfsEvent::TreeEdge(u, v, weight) => { + let edge = (u.index(), v.index(), weight); + vis.tree_edge.call1(py, (edge,)) + } + DfsEvent::BackEdge(u, v, weight) => { + let edge = (u.index(), v.index(), weight); + vis.back_edge.call1(py, (edge,)) + } + DfsEvent::CrossForwardEdge(u, v, weight) => { + let edge = (u.index(), v.index(), weight); + vis.forward_or_cross_edge.call1(py, (edge,)) + } + DfsEvent::Finish(u, Time(t)) => { + vis.finish_vertex.call1(py, (u.index(), t)) + } + }; + + match res { + Err(e) => { + if e.is_instance::(py) { + Ok(Control::Prune) + } else { + Err(e) + } + } + Ok(_) => Ok(Control::Continue), + } +} diff --git a/src/traversal/mod.rs b/src/traversal/mod.rs index c0a9edac1..f28ba6f69 100644 --- a/src/traversal/mod.rs +++ b/src/traversal/mod.rs @@ -11,9 +11,14 @@ // under the License. mod bfs_visit; +pub mod dfs_visit; use bfs_visit::{bfs_handler, PyBfsVisitor}; -use retworkx_core::traversal::{breadth_first_search, dfs_edges}; +use dfs_visit::{dfs_handler, PyDfsVisitor}; + +use retworkx_core::traversal::{ + breadth_first_search, depth_first_search, dfs_edges, +}; use super::{digraph, graph, iterators}; @@ -385,3 +390,171 @@ pub fn graph_bfs_search( Ok(()) } + +/// Depth-first traversal of a directed graph. +/// +/// The pseudo-code for the DFS algorithm is listed below, with the annotated +/// event points, for which the given visitor object will be called with the +/// appropriate method. +/// +/// :: +/// +/// DFS(G) +/// for each vertex u in V +/// color[u] := WHITE initialize vertex u +/// end for +/// time := 0 +/// call DFS-VISIT(G, source) start vertex s +/// +/// DFS-VISIT(G, u) +/// color[u] := GRAY discover vertex u +/// for each v in Adj[u] examine edge (u,v) +/// if (color[v] = WHITE) (u,v) is a tree edge +/// all DFS-VISIT(G, v) +/// else if (color[v] = GRAY) (u,v) is a back edge +/// ... +/// else if (color[v] = BLACK) (u,v) is a cross or forward edge +/// ... +/// end for +/// color[u] := BLACK finish vertex u +/// +/// 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. 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 DFSVisitor +/// +/// class TreeEdgesRecorder(DFSVisitor): +/// +/// 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.dfs_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 depth-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.DFSVisitor`. +#[pyfunction] +#[pyo3(text_signature = "(graph, source, visitor)")] +pub fn digraph_dfs_search( + py: Python, + graph: &digraph::PyDiGraph, + source: Option>, + visitor: PyDfsVisitor, +) -> PyResult<()> { + let starts: Vec<_> = match source { + Some(nx) => nx.into_iter().map(NodeIndex::new).collect(), + None => graph.graph.node_indices().collect(), + }; + + depth_first_search(&graph.graph, starts, |event| { + dfs_handler(py, &visitor, event) + })?; + + Ok(()) +} + +/// Depth-first traversal of an undirected graph. +/// +/// The pseudo-code for the DFS algorithm is listed below, with the annotated +/// event points, for which the given visitor object will be called with the +/// appropriate method. +/// +/// :: +/// +/// DFS(G) +/// for each vertex u in V +/// color[u] := WHITE initialize vertex u +/// end for +/// time := 0 +/// call DFS-VISIT(G, source) start vertex s +/// +/// DFS-VISIT(G, u) +/// color[u] := GRAY discover vertex u +/// for each v in Adj[u] examine edge (u,v) +/// if (color[v] = WHITE) (u,v) is a tree edge +/// all DFS-VISIT(G, v) +/// else if (color[v] = GRAY) (u,v) is a back edge +/// ... +/// else if (color[v] = BLACK) (u,v) is a cross or forward edge +/// ... +/// end for +/// color[u] := BLACK finish vertex u +/// +/// 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. 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 DFSVisitor +/// +/// class TreeEdgesRecorder(DFSVisitor): +/// +/// 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.dfs_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 depth-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.DFSVisitor`. +#[pyfunction] +#[pyo3(text_signature = "(graph, source, visitor)")] +pub fn graph_dfs_search( + py: Python, + graph: &graph::PyGraph, + source: Option>, + visitor: PyDfsVisitor, +) -> PyResult<()> { + let starts: Vec<_> = match source { + Some(nx) => nx.into_iter().map(NodeIndex::new).collect(), + None => graph.graph.node_indices().collect(), + }; + + depth_first_search(&graph.graph, starts, |event| { + dfs_handler(py, &visitor, event) + })?; + + Ok(()) +} diff --git a/tests/digraph/test_dfs_search.py b/tests/digraph/test_dfs_search.py new file mode 100644 index 000000000..f582e2181 --- /dev/null +++ b/tests/digraph/test_dfs_search.py @@ -0,0 +1,108 @@ +# 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 TestDfsSearch(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_dfs_tree_edges(self): + class TreeEdgesRecorder(retworkx.visit.DFSVisitor): + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append((edge[0], edge[1])) + + vis = TreeEdgesRecorder() + retworkx.digraph_dfs_search(self.graph, [0], vis) + self.assertEqual(vis.edges, [(0, 2), (2, 6), (2, 5), (5, 3), (2, 1)]) + + def test_digraph_dfs_tree_edges_no_starting_point(self): + class TreeEdgesRecorder(retworkx.visit.DFSVisitor): + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append((edge[0], edge[1])) + + vis = TreeEdgesRecorder() + retworkx.digraph_dfs_search(self.graph, None, vis) + self.assertEqual( + vis.edges, [(0, 2), (2, 6), (2, 5), (5, 3), (2, 1), (4, 7)] + ) + + def test_digraph_dfs_tree_edges_restricted(self): + class TreeEdgesRecorderRestricted(retworkx.visit.DFSVisitor): + + prohibited = [(0, 1), (5, 3)] + + 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_dfs_search(self.graph, [0], vis) + self.assertEqual(vis.edges, [(0, 2), (2, 6), (2, 5), (2, 1), (1, 3)]) + + def test_digraph_dfs_goal_search(self): + class GoalSearch(retworkx.visit.DFSVisitor): + + 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() + try: + retworkx.digraph_dfs_search(self.graph, [0], vis) + except retworkx.visit.StopSearch: + pass + self.assertEqual(vis.reconstruct_path(), [0, 2, 5, 3]) diff --git a/tests/graph/test_dfs_search.py b/tests/graph/test_dfs_search.py new file mode 100644 index 000000000..6926e865f --- /dev/null +++ b/tests/graph/test_dfs_search.py @@ -0,0 +1,108 @@ +# 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 TestDfsSearch(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_dfs_tree_edges(self): + class TreeEdgesRecorder(retworkx.visit.DFSVisitor): + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append((edge[0], edge[1])) + + vis = TreeEdgesRecorder() + retworkx.graph_dfs_search(self.graph, [0], vis) + self.assertEqual(vis.edges, [(0, 2), (2, 6), (2, 5), (5, 3), (3, 1)]) + + def test_graph_dfs_tree_edges_no_starting_point(self): + class TreeEdgesRecorder(retworkx.visit.DFSVisitor): + def __init__(self): + self.edges = [] + + def tree_edge(self, edge): + self.edges.append((edge[0], edge[1])) + + vis = TreeEdgesRecorder() + retworkx.graph_dfs_search(self.graph, None, vis) + self.assertEqual( + vis.edges, [(0, 2), (2, 6), (2, 5), (5, 3), (3, 1), (4, 7)] + ) + + def test_graph_dfs_tree_edges_restricted(self): + class TreeEdgesRecorderRestricted(retworkx.visit.DFSVisitor): + + 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_dfs_search(self.graph, [0], vis) + self.assertEqual(vis.edges, [(0, 1), (1, 3), (3, 5), (5, 2), (2, 6)]) + + def test_graph_dfs_goal_search(self): + class GoalSearch(retworkx.visit.DFSVisitor): + + 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() + try: + retworkx.graph_dfs_search(self.graph, [0], vis) + except retworkx.visit.StopSearch: + pass + self.assertEqual(vis.reconstruct_path(), [0, 2, 5, 3])