retworkx_core/traversal/
dfs_edges.rs

1// Licensed under the Apache License, Version 2.0 (the "License"); you may
2// not use this file except in compliance with the License. You may obtain
3// a copy of the License at
4//
5//     http://www.apache.org/licenses/LICENSE-2.0
6//
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10// License for the specific language governing permissions and limitations
11// under the License.
12
13use std::hash::Hash;
14
15use hashbrown::{HashMap, HashSet};
16
17use petgraph::visit::{
18    EdgeCount, IntoNeighbors, IntoNodeIdentifiers, NodeCount, NodeIndexable,
19    Visitable,
20};
21
22/// Return an edge list of the tree edges from a depth-first traversal.
23///
24/// The pseudo-code for the DFS algorithm is listed below. The output
25/// contains the tree edges found by the procedure.
26///
27/// ```norust
28/// DFS(G, v)
29///   let S be a stack
30///   label v as discovered
31///   PUSH(S, (v, iterator of G.neighbors(v)))
32///   while (S != Ø)
33///       let (v, iterator) := LAST(S)
34///       if hasNext(iterator) then
35///           w := next(iterator)
36///           if w is not labeled as discovered then
37///               label w as discovered                   # (v, w) is a tree edge
38///               PUSH(S, (w, iterator of G.neighbors(w)))
39///       else
40///           POP(S)
41///   end while
42/// ```
43///
44/// Arguments:
45///
46/// * `graph` - the graph to run on
47/// * `source` - the optional node index to use as the starting node for the
48///     depth-first search. If specified the edge list will only return edges
49///     in the components reachable from this index. If this is not specified
50///     then a source will be chosen arbitrarily and repeated until all
51///     components of the graph are searched
52///
53/// # Example
54/// ```rust
55/// use retworkx_core::petgraph;
56/// use retworkx_core::traversal::dfs_edges;
57///
58/// let g = petgraph::graph::UnGraph::<i32, ()>::from_edges(&[
59///     (0, 1), (1, 2), (1, 3), (2, 4), (3, 4)
60/// ]);
61/// let dfs_edges = dfs_edges(&g, Some(petgraph::graph::NodeIndex::new(0)));
62/// assert_eq!(vec![(0, 1), (1, 2), (2, 4), (4, 3)], dfs_edges);
63/// ```
64pub fn dfs_edges<G>(graph: G, source: Option<G::NodeId>) -> Vec<(usize, usize)>
65where
66    G: IntoNodeIdentifiers
67        + NodeIndexable
68        + IntoNeighbors
69        + NodeCount
70        + EdgeCount
71        + Visitable,
72    G::NodeId: Eq + Hash,
73{
74    let nodes: Vec<G::NodeId> = match source {
75        Some(start) => vec![start],
76        None => graph.node_identifiers().collect(),
77    };
78    let node_count = graph.node_count();
79    let mut visited: HashSet<G::NodeId> = HashSet::with_capacity(node_count);
80    // Avoid potential overallocation if source node is provided
81    let mut out_vec: Vec<(usize, usize)> = if source.is_some() {
82        Vec::new()
83    } else {
84        Vec::with_capacity(core::cmp::min(
85            graph.node_count() - 1,
86            graph.edge_count(),
87        ))
88    };
89    for start in nodes {
90        if visited.contains(&start) {
91            continue;
92        }
93        visited.insert(start);
94        let mut children: Vec<G::NodeId> = graph.neighbors(start).collect();
95        children.reverse();
96        let mut stack: Vec<(G::NodeId, Vec<G::NodeId>)> =
97            vec![(start, children)];
98        // Used to track the last position in children vec across iterations
99        let mut index_map: HashMap<G::NodeId, usize> =
100            HashMap::with_capacity(node_count);
101        index_map.insert(start, 0);
102        while !stack.is_empty() {
103            let temp_parent = stack.last().unwrap();
104            let parent = temp_parent.0;
105            let children = temp_parent.1.clone();
106            let count = *index_map.get(&parent).unwrap();
107            let mut found = false;
108            let mut index = count;
109            for child in &children[index..] {
110                index += 1;
111                if !visited.contains(child) {
112                    out_vec
113                        .push((graph.to_index(parent), graph.to_index(*child)));
114                    visited.insert(*child);
115                    let mut grandchildren: Vec<G::NodeId> =
116                        graph.neighbors(*child).collect();
117                    grandchildren.reverse();
118                    stack.push((*child, grandchildren));
119                    index_map.insert(*child, 0);
120                    *index_map.get_mut(&parent).unwrap() = index;
121                    found = true;
122                    break;
123                }
124            }
125            if !found || children.is_empty() {
126                stack.pop();
127            }
128        }
129    }
130    out_vec
131}
retworkx_core/traversal/dfs_edges.rs

retworkx_core/traversal/
dfs_edges.rs
