rustworkx_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, Visitable,
19};
20
21/// Return an edge list of the tree edges from a depth-first traversal.
22///
23/// The pseudo-code for the DFS algorithm is listed below. The output
24/// contains the tree edges found by the procedure.
25///
26/// ```norust
27/// DFS(G, v)
28///   let S be a stack
29///   label v as discovered
30///   PUSH(S, (v, iterator of G.neighbors(v)))
31///   while (S != Ø)
32///       let (v, iterator) := LAST(S)
33///       if hasNext(iterator) then
34///           w := next(iterator)
35///           if w is not labeled as discovered then
36///               label w as discovered                   # (v, w) is a tree edge
37///               PUSH(S, (w, iterator of G.neighbors(w)))
38///       else
39///           POP(S)
40///   end while
41/// ```
42///
43/// Arguments:
44///
45/// * `graph` - the graph to run on
46/// * `source` - the optional node index to use as the starting node for the
47///   depth-first search. If specified the edge list will only return edges
48///   in the components reachable from this index. If this is not specified
49///   then a source will be chosen arbitrarily and repeated until all
50///   components of the graph are searched
51///
52/// # Example
53/// ```rust
54/// use rustworkx_core::petgraph;
55/// use rustworkx_core::traversal::dfs_edges;
56///
57/// let g = petgraph::graph::UnGraph::<i32, ()>::from_edges(&[
58///     (0, 1), (1, 2), (1, 3), (2, 4), (3, 4)
59/// ]);
60/// let dfs_edges = dfs_edges(&g, Some(petgraph::graph::NodeIndex::new(0)));
61/// assert_eq!(vec![(0, 1), (1, 2), (2, 4), (4, 3)], dfs_edges);
62/// ```
63pub fn dfs_edges<G>(graph: G, source: Option<G::NodeId>) -> Vec<(usize, usize)>
64where
65    G: IntoNodeIdentifiers + NodeIndexable + IntoNeighbors + NodeCount + EdgeCount + Visitable,
66    G::NodeId: Eq + Hash,
67{
68    let nodes: Vec<G::NodeId> = match source {
69        Some(start) => vec![start],
70        None => graph.node_identifiers().collect(),
71    };
72    let node_count = graph.node_count();
73    let mut visited: HashSet<G::NodeId> = HashSet::with_capacity(node_count);
74    // Avoid potential overallocation if source node is provided
75    let mut out_vec: Vec<(usize, usize)> = if source.is_some() {
76        Vec::new()
77    } else {
78        Vec::with_capacity(core::cmp::min(graph.node_count() - 1, graph.edge_count()))
79    };
80    for start in nodes {
81        if visited.contains(&start) {
82            continue;
83        }
84        visited.insert(start);
85        let mut children: Vec<G::NodeId> = graph.neighbors(start).collect();
86        children.reverse();
87        let mut stack: Vec<(G::NodeId, Vec<G::NodeId>)> = vec![(start, children)];
88        // Used to track the last position in children vec across iterations
89        let mut index_map: HashMap<G::NodeId, usize> = HashMap::with_capacity(node_count);
90        index_map.insert(start, 0);
91        while !stack.is_empty() {
92            let temp_parent = stack.last().unwrap();
93            let parent = temp_parent.0;
94            let children = temp_parent.1.clone();
95            let count = *index_map.get(&parent).unwrap();
96            let mut found = false;
97            let mut index = count;
98            for child in &children[index..] {
99                index += 1;
100                if !visited.contains(child) {
101                    out_vec.push((graph.to_index(parent), graph.to_index(*child)));
102                    visited.insert(*child);
103                    let mut grandchildren: Vec<G::NodeId> = graph.neighbors(*child).collect();
104                    grandchildren.reverse();
105                    stack.push((*child, grandchildren));
106                    index_map.insert(*child, 0);
107                    *index_map.get_mut(&parent).unwrap() = index;
108                    found = true;
109                    break;
110                }
111            }
112            if !found || children.is_empty() {
113                stack.pop();
114            }
115        }
116    }
117    out_vec
118}
rustworkx_core/traversal/dfs_edges.rs

rustworkx_core/traversal/
dfs_edges.rs
