From 89277f044b3e3eb1e8bea593a215d4a1333b6f23 Mon Sep 17 00:00:00 2001 From: georgios-ts <45130028+georgios-ts@users.noreply.github.com> Date: Mon, 10 Jan 2022 23:46:11 +0200 Subject: [PATCH] Add dfs-search (#456) * Add dfs-search Widely used graph libraries like Boost Graph Library and graph-tool provide the ability to insert callback functions at specified event points for many common graph search algorithms. petgraph has a similar concept in `petgraph::visit::depth_first_search` function. This commit implements an iterative version of `depth_first_search`, that will be used in a follow-up in the pending PRs #444, #445. At the same time it exposes this new functionality in Python by letting users subclassing `retworkx.visit.DFSVisitor` and provide their own implementation for the appropriate callback functions. The benefit is less memory consumption since we avoid storing the results but rather let the user take the desired action at specified points. For example, if a user wants to process the nodes in dfs-order, we don't need to create a new list with all the graph nodes in dfs-order but rather the user can process a node on the fly. We can (probably) leverage this approach in other algorithms as an alternative for our inability to provide "real" python iterators. * break or prune the search tree with custom exceptions * tests * move `depth_first_search` to retworkx-core + create a new traversal module together with `dfs_edges` function * add option to pass a vector of starting nodes + in dfs events report the weight of an edge * lint * fix docs * this is the last fix * define custom exceptions `StopSearch`, `PruneSearch` in python and import them in rust so we can simplify a bit the namespaces * minor doc fixes * ignore flake warning * run black * deduplicate macro and mention petgraph DfsEvent struct --- docs/source/api.rst | 4 + .../notes/dfs-search-6083680bf62356b0.yaml | 28 ++ retworkx-core/src/traversal/bfs_visit.rs | 22 +- retworkx-core/src/traversal/dfs_visit.rs | 260 ++++++++++++++++++ retworkx-core/src/traversal/mod.rs | 26 ++ retworkx/__init__.py | 80 ++++++ retworkx/visit.py | 51 ++++ src/lib.rs | 2 + src/traversal/dfs_visit.rs | 66 +++++ src/traversal/mod.rs | 175 +++++++++++- tests/digraph/test_dfs_search.py | 108 ++++++++ tests/graph/test_dfs_search.py | 108 ++++++++ 12 files changed, 908 insertions(+), 22 deletions(-) create mode 100644 releasenotes/notes/dfs-search-6083680bf62356b0.yaml create mode 100644 retworkx-core/src/traversal/dfs_visit.rs create mode 100644 src/traversal/dfs_visit.rs create mode 100644 tests/digraph/test_dfs_search.py create mode 100644 tests/graph/test_dfs_search.py 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])