leetcode/util/
tree_node.rs

1use std::rc::Rc;
2use std::cell::RefCell;
3
4/// LeetCode 题目中用到的树的节点结构
5// Definition for a binary tree node.
6#[derive(Debug, PartialEq, Eq)]
7pub struct TreeNode {
8    pub val: i32,
9    pub left: Option<Rc<RefCell<TreeNode>>>,
10    pub right: Option<Rc<RefCell<TreeNode>>>,
11}
12
13impl TreeNode {
14    #[inline]
15    pub fn new(val: i32) -> Self {
16        TreeNode {
17            val,
18            left: None,
19            right: None
20        }
21    }
22
23    /// 从 `Vec<Option<i32>>` 生成树结构。一般建议使用宏 [`tree!`](tree)
24    pub fn from_vec(vec: Vec<Option<i32>>) -> Option<Rc<RefCell<TreeNode>>> {
25        use std::collections::VecDeque;
26        let root = Some(Rc::new(RefCell::new(TreeNode::new(vec[0].unwrap()))));
27        let mut queue = VecDeque::new();
28        queue.push_back(root.as_ref().unwrap().clone());
29
30        for children in vec[1..].chunks(2) {
31            let parent = queue.pop_front().unwrap();
32            if let Some(v) = children[0] {
33                parent.borrow_mut().left = Some(Rc::new(RefCell::new(TreeNode::new(v))));
34                queue.push_back(parent.borrow().left.as_ref().unwrap().clone());
35            }
36            if children.len() > 1 {
37                if let Some(v) = children[1] {
38                    parent.borrow_mut().right = Some(Rc::new(RefCell::new(TreeNode::new(v))));
39                    queue.push_back(parent.borrow().right.as_ref().unwrap().clone());
40                }
41            }
42        }
43        root
44    }
45}
46
47/// Generate a tree structure from a vec-like syntax.
48///
49/// For example:
50///
51/// ```no_run
52/// #[macro_use] extern crate leetcode;
53/// use std::rc::Rc;
54/// use std::cell::RefCell;
55/// use leetcode::{tree, TreeNode};
56///
57/// let tree_node: Option<Rc<RefCell<TreeNode>>> = tree![4,2,7,1,3,6,9];
58/// ```
59/// will be expanded to an optional value: `Option<Rc<RefCell<TreeNode>>>` which with the following
60/// tree structure:
61///
62/// ```text
63///      4
64///    /   \
65///   2     7
66///  / \   / \
67/// 1   3 6   9
68/// ```
69///
70/// And:
71///
72/// ```no_run
73/// #[macro_use] extern crate leetcode;
74/// use std::rc::Rc;
75/// use std::cell::RefCell;
76/// use leetcode::{tree, TreeNode};
77///
78/// let tree_node: Option<Rc<RefCell<TreeNode>>> = tree![4,2,7,null,null,6,9];
79/// ```
80///
81/// will be expanded to an optional value: `Option<Rc<RefCell<TreeNode>>>` which with the following
82/// tree structure:
83///
84/// ```text
85///        4
86///      /   \
87///     2     7
88///    / \   / \
89///         6   9
90/// ```
91#[macro_export]
92macro_rules! tree {
93    () => {
94        None
95    };
96    ($($e:expr),*) => {
97        {
98            let vec = vec![$(stringify!($e)), *];
99            let vec = vec.into_iter().map(|v| v.parse::<i32>().ok()).collect::<Vec<_>>();
100            TreeNode::from_vec(vec)
101        }
102    };
103    ($($e:expr,)*) => {(tree![$($e),*])};
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn test_from_vec() {
112        let tree = tree![4,2,7,1,3,6,9];
113        println!("{:?}", tree);
114        assert_eq!(None::<Option<Rc<RefCell<TreeNode>>>>, tree![]);
115    }
116}
leetcode/util/tree_node.rs

leetcode/util/
tree_node.rs
