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>;