diff --git a/src/lib.rs b/src/lib.rs index a98c81e..ee482f9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -96,6 +96,26 @@ impl Graph { return Ok(leaf_keys); } + + /// Return all keys ordered like you were doing + /// a depth first search. + pub fn get_keys_by_depth(&self) -> Result, GraphError> { + let mut visited_keys = Vec::::new(); + let mut stack = Vec::::new(); + + stack.push(self.root_key()); + + while stack.len() > 0 { + let next_key = stack.pop().unwrap(); + if !visited_keys.contains(&next_key) { + visited_keys.push(next_key); + let children_of_next_key = self.children_of(&next_key)?; + stack.extend(children_of_next_key); + } + } + + Ok(visited_keys) + } } #[derive(Debug, Clone)] @@ -178,7 +198,15 @@ mod tests { let fourth_add_res = graph.add(fourth_val, 1); assert!(fourth_add_res.is_ok()); + let fifth_val = NodeType::new(5.0); + let fifth_add_res = graph.add(fifth_val, 2); + assert!(fifth_add_res.is_ok()); + let leaf_keys = graph.find_leaf_keys().unwrap(); assert!(leaf_keys.len() == 2); + + let keys_by_depth = graph.get_keys_by_depth().unwrap(); + let expected_keys_by_depth = vec![0, 2, 4, 1, 3]; + assert!(keys_by_depth == expected_keys_by_depth); } }