causal_hub/graphs/algorithms/traversal/
depth_first_search.rs

1use std::{
2    collections::{hash_map::Entry, HashMap, VecDeque},
3    vec::Vec,
4};
5
6use super::Traversal;
7use crate::{
8    graphs::{directions, BaseGraph, DirectedGraph, UndirectedGraph},
9    Ch, Ne, V,
10};
11
12/// Depth-first search structure.
13///
14/// This structure contains the `discovery_time`, `finish_time` and `predecessor` maps.
15///
16pub struct DepthFirstSearch<'a, G, D>
17where
18    G: BaseGraph<Direction = D>,
19{
20    /// Given graph reference.
21    g: &'a G,
22    /// The visit stack.
23    stack: Vec<usize>,
24    /// Global time counter.
25    pub time: usize,
26    /// Discovery time of each discovered vertex.
27    pub discovery_time: HashMap<usize, usize>,
28    /// Finish time of each discovered vertex.
29    pub finish_time: HashMap<usize, usize>,
30    /// Predecessor of each discovered vertex (except the source vertex).
31    pub predecessor: HashMap<usize, usize>,
32}
33
34impl<'a, G, D> DepthFirstSearch<'a, G, D>
35where
36    G: BaseGraph<Direction = D>,
37{
38    /// Build a DFS iterator.
39    ///
40    /// Build a DFS iterator for a given graph. This will execute the [`Tree`](super::Traversal)
41    /// variant of the algorithm, if not specified otherwise.
42    ///
43    /// # Panics
44    ///
45    /// Panics if the (optional) source vertex is not in the graph.
46    ///
47    /// # Examples
48    ///
49    /// ```
50    /// use causal_hub::prelude::*;
51    ///
52    /// // Define edge set.
53    /// let e = EdgeList::from([
54    ///     ("A", "B"), ("B", "C"), ("A", "D"), ("A", "E"), ("E", "F")
55    /// ]);
56    ///
57    /// // Build a new graph.
58    /// let mut g = DiGraph::from(e);
59    ///
60    /// // Build the search object over said graph with `0` as source vertex.
61    /// let mut search = DFS::from((&g, 0));
62    ///
63    /// // Collect visit order by reference, preserving search structure.
64    /// let order: Vec<_> = search.by_ref().collect();
65    /// // The visit returns a pre-order.
66    /// assert_eq!(order, [0, 1, 2, 3, 4, 5]);
67    ///
68    /// // The source vertex has discovery-time equals to zero ...
69    /// assert_eq!(search.discovery_time[&0], 0);
70    /// // ... finish-time equals to two times the number of discovered vertices minus one ...
71    /// assert_eq!(search.finish_time[&0], 2 * search.discovery_time.len() - 1);
72    /// // ... and no predecessor by definition.
73    /// assert_eq!(search.predecessor.contains_key(&0), false);
74    ///
75    /// // For example, vertex `5` has discovery-time equals to eight ...
76    /// assert_eq!(search.discovery_time[&5], 8);
77    /// // ... finish-time equals to nine ...
78    /// assert_eq!(search.finish_time[&5], 9);
79    /// // ... and its predecessor is `4`.
80    /// assert_eq!(search.predecessor[&5], 4);
81    /// ```
82    ///
83    pub fn new(g: &'a G, x: Option<usize>, m: Traversal) -> Self {
84        // Initialize default search object.
85        let mut search = Self {
86            // Set target graph.
87            g,
88            // Initialize the to-be-visited queue with the source vertex.
89            stack: Default::default(),
90            // Initialize the global clock.
91            time: Default::default(),
92            // Initialize the discovery-time map.
93            discovery_time: Default::default(),
94            // Initialize the finish-time map.
95            finish_time: Default::default(),
96            // Initialize the predecessor map.
97            predecessor: Default::default(),
98        };
99        // If the graph is null.
100        if g.order() == 0 {
101            // Assert source vertex is none.
102            assert!(x.is_none());
103            // Then, return the default search object.
104            return search;
105        }
106        // Get source vertex, if any.
107        let x = match x {
108            // If no source vertex is given, choose the first one in the vertex set.
109            None => V!(g).next().unwrap(),
110            // Otherwise ...
111            Some(x) => {
112                // ... assert that source vertex is in graph.
113                assert!(g.has_vertex(x));
114                // Return given source vertex.
115                x
116            }
117        };
118        // If visit variant is Forest.
119        if matches!(m, Traversal::Forest) {
120            // Add vertices to the visit stack in reverse to preserve order.
121            let mut queue = VecDeque::<usize>::with_capacity(g.order());
122            queue.extend(V!(g).filter(|&y| y != x));
123            search.stack.extend(queue.iter().rev());
124        }
125        // Push source vertex onto the stack.
126        search.stack.push(x);
127        // Return search object.
128        search
129    }
130}
131
132impl<'a, G> Iterator for DepthFirstSearch<'a, G, directions::Undirected>
133where
134    G: UndirectedGraph<Direction = directions::Undirected>,
135{
136    type Item = usize;
137
138    fn next(&mut self) -> Option<Self::Item> {
139        // If there are still vertices to be visited.
140        while let Some(&x) = self.stack.last() {
141            // Check if vertex is WHITE (i.e. was not seen before).
142            if let Entry::Vacant(e) = self.discovery_time.entry(x) {
143                // Set its discover time (as GRAY).
144                e.insert(self.time);
145                // Increment time.
146                self.time += 1;
147                // Initialize visiting queue.
148                let mut queue = VecDeque::new();
149                // Iterate over reachable vertices.
150                for y in Ne!(self.g, x) {
151                    // Filter already visited vertices (as GRAY).
152                    if !self.discovery_time.contains_key(&y) {
153                        // Set predecessor.
154                        self.predecessor.insert(y, x);
155                        // Add to queue.
156                        queue.push_front(y);
157                    }
158                }
159                // Push vertices onto the stack in reverse order, this makes
160                // traversal order and neighborhood order the same.
161                self.stack.extend(queue);
162                // Return vertex in pre-order.
163                return Some(x);
164            // If the vertex is NOT WHITE.
165            } else {
166                // Remove it from stack.
167                self.stack.pop();
168                // Check if it is GRAY (not BLACK).
169                if let Entry::Vacant(e) = self.finish_time.entry(x) {
170                    // Set its finish time (as BLACK).
171                    e.insert(self.time);
172                    // Increment time.
173                    self.time += 1;
174                }
175            }
176        }
177
178        None
179    }
180}
181
182impl<'a, G> Iterator for DepthFirstSearch<'a, G, directions::Directed>
183where
184    G: DirectedGraph<Direction = directions::Directed>,
185{
186    type Item = usize;
187
188    fn next(&mut self) -> Option<Self::Item> {
189        // If there are still vertices to be visited.
190        while let Some(&x) = self.stack.last() {
191            // Check if vertex is WHITE (i.e. was not seen before).
192            if let Entry::Vacant(e) = self.discovery_time.entry(x) {
193                // Set its discover time (as GRAY).
194                e.insert(self.time);
195                // Increment time.
196                self.time += 1;
197                // Initialize visiting queue.
198                let mut queue = VecDeque::new();
199                // Iterate over reachable vertices.
200                for y in Ch!(self.g, x) {
201                    // Filter already visited vertices (as GRAY).
202                    if !self.discovery_time.contains_key(&y) {
203                        // Set predecessor.
204                        self.predecessor.insert(y, x);
205                        // Add to queue.
206                        queue.push_front(y);
207                    }
208                }
209                // Push vertices onto the stack in reverse order, this makes
210                // traversal order and neighborhood order the same.
211                self.stack.extend(queue);
212                // Return vertex in pre-order.
213                return Some(x);
214            // If the vertex is NOT WHITE.
215            } else {
216                // Remove it from stack.
217                self.stack.pop();
218                // Check if it is GRAY (not BLACK).
219                if let Entry::Vacant(e) = self.finish_time.entry(x) {
220                    // Set its finish time (as BLACK).
221                    e.insert(self.time);
222                    // Increment time.
223                    self.time += 1;
224                }
225            }
226        }
227
228        None
229    }
230}
231
232impl<'a, G, D> From<&'a G> for DepthFirstSearch<'a, G, D>
233where
234    G: BaseGraph<Direction = D>,
235{
236    fn from(g: &'a G) -> Self {
237        Self::new(g, None, Traversal::Tree)
238    }
239}
240
241impl<'a, G, D> From<(&'a G, usize)> for DepthFirstSearch<'a, G, D>
242where
243    G: BaseGraph<Direction = D>,
244{
245    fn from((g, x): (&'a G, usize)) -> Self {
246        Self::new(g, Some(x), Traversal::Tree)
247    }
248}
249
250/// Alias for depth-first search.
251pub type DFS<'a, G, D> = DepthFirstSearch<'a, G, D>;