leetcode_rust/
binary_tree_level_order_traversal.rs

1#![allow(dead_code)]
2
3use std::cell::RefCell;
4use std::collections::VecDeque;
5use std::rc::Rc;
6
7use crate::binary_tree::TreeNode;
8
9pub fn level_order(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<Vec<i32>> {
10    let mut res: Vec<Vec<i32>> = vec![];
11    let mut q = VecDeque::new();
12    q.push_back((root, 0));
13
14    while !q.is_empty() {
15        let cur = q.pop_front().unwrap();
16        let depth = cur.1;
17
18        let cur = cur.0;
19        match cur {
20            Some(node) => {
21                if depth >= res.len() {
22                    res.push(vec![]);
23                }
24                res[depth].push(node.borrow().val);
25                q.push_back((node.borrow().left.clone(), depth + 1));
26                q.push_back((node.borrow().right.clone(), depth + 1));
27            }
28            _ => {}
29        }
30    }
31
32    res
33}
34
35#[cfg(test)]
36mod tests {
37    use super::*;
38
39    #[test]
40    fn test1() {
41        let tree = Some(Rc::new(RefCell::new(TreeNode {
42            val: 1,
43            left: None,
44            right: Some(Rc::new(RefCell::new(TreeNode {
45                val: 2,
46                left: Some(Rc::new(RefCell::new(TreeNode {
47                    val: 3,
48                    left: None,
49                    right: None,
50                }))),
51                right: None,
52            }))),
53        })));
54        let res = level_order(tree);
55        assert_eq!(res, vec![vec![1], vec![2], vec![3]]);
56    }
57}