// At Level โ 99 Problems #33
// Collect the nodes at a given level of a binary tree.
// Root is at level 1.
#[derive(Debug, PartialEq, Clone)]
enum Tree<T> {
Leaf,
Node(Box<Tree<T>>, T, Box<Tree<T>>),
}
impl<T: Clone> Tree<T> {
fn leaf() -> Self { Tree::Leaf }
fn node(l: Tree<T>, v: T, r: Tree<T>) -> Self {
Tree::Node(Box::new(l), v, Box::new(r))
}
}
fn at_level<T: Clone>(tree: &Tree<T>, level: usize) -> Vec<T> {
match tree {
Tree::Leaf => vec![],
Tree::Node(l, v, r) => {
if level == 1 {
vec![v.clone()]
} else {
let mut result = at_level(l, level - 1);
result.extend(at_level(r, level - 1));
result
}
}
}
}
/// Breadth-first traversal returning all levels.
fn levels<T: Clone>(tree: &Tree<T>) -> Vec<Vec<T>> {
let mut result = Vec::new();
let mut queue = vec![tree];
while !queue.is_empty() {
let mut next_queue = Vec::new();
let mut level_vals = Vec::new();
for t in queue {
if let Tree::Node(l, v, r) = t {
level_vals.push(v.clone());
if !matches!(l.as_ref(), Tree::Leaf) {
next_queue.push(l.as_ref());
}
if !matches!(r.as_ref(), Tree::Leaf) {
next_queue.push(r.as_ref());
}
}
}
if !level_vals.is_empty() {
result.push(level_vals);
}
queue = next_queue;
}
result
}
fn sample_tree() -> Tree<char> {
// a level 1
// / \
// b c level 2
// / \
// d e level 3
Tree::node(
Tree::node(
Tree::node(Tree::leaf(), 'd', Tree::leaf()),
'b',
Tree::node(Tree::leaf(), 'e', Tree::leaf()),
),
'a',
Tree::node(Tree::leaf(), 'c', Tree::leaf()),
)
}
fn main() {
let t = sample_tree();
println!("Level 1: {:?}", at_level(&t, 1)); // [a]
println!("Level 2: {:?}", at_level(&t, 2)); // [b, c]
println!("Level 3: {:?}", at_level(&t, 3)); // [d, e]
println!("Level 4: {:?}", at_level(&t, 4)); // []
println!("All levels: {:?}", levels(&t));
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_level_1() {
assert_eq!(at_level(&sample_tree(), 1), vec!['a']);
}
#[test]
fn test_level_2() {
assert_eq!(at_level(&sample_tree(), 2), vec!['b', 'c']);
}
#[test]
fn test_level_3() {
assert_eq!(at_level(&sample_tree(), 3), vec!['d', 'e']);
}
#[test]
fn test_level_beyond() {
assert_eq!(at_level(&sample_tree(), 5), vec![]);
}
#[test]
fn test_all_levels_count() {
let all = levels(&sample_tree());
assert_eq!(all.len(), 3);
assert_eq!(all[0], vec!['a']);
assert_eq!(all[1], vec!['b', 'c']);
assert_eq!(all[2], vec!['d', 'e']);
}
}
(* At Level *)
(* OCaml 99 Problems #33 *)
(* Implementation for example 33 *)
(* Tests *)
let () =
(* Add tests *)
print_endline "โ OCaml tests passed"