leetcode/tree/medium/count_nodes/src/lib.rs at master · zhangyuang/leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
/*
 * @lc app=leetcode.cn id=222 lang=rust
 *
 * [222] 完全二叉树的节点个数
 */
//Definition for a binary tree node.
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
    pub val: i32,
    pub left: Option<Rc<RefCell<TreeNode>>>,
    pub right: Option<Rc<RefCell<TreeNode>>>,
}

impl TreeNode {
    #[inline]
    pub fn new(val: i32) -> Self {
        TreeNode {
            val,
            left: None,
            right: None,
        }
    }
}
struct Solution {}
// @lc code=start
// Definition for a binary tree node.
// #[derive(Debug, PartialEq, Eq)]
// pub struct TreeNode {
//   pub val: i32,
//   pub left: Option<Rc<RefCell<TreeNode>>>,
//   pub right: Option<Rc<RefCell<TreeNode>>>,
// }
//
// impl TreeNode {
//   #[inline]
//   pub fn new(val: i32) -> Self {
//     TreeNode {
//       val,
//       left: None,
//       right: None
//     }
//   }
// }
use std::cell::RefCell;
use std::cmp::max;
use std::collections::VecDeque;
use std::rc::Rc;

impl Solution {
    // 解法1，计算高度，加上叶子节点的个数
    // pub fn count_nodes(root: Option<Rc<RefCell<TreeNode>>>) -> i32 {
    //     let height = Self::get_height(&root).unwrap();
    //     if height == 1 || height == 0 {
    //         return height;
    //     }
    //     let leaf = Self::get_leaf(&root, 1, height).unwrap();
    //     let two: i32 = 2;
    //     return two.pow((height - 1) as u32) - 1 + leaf;
    // }
    // fn get_leaf(
    //     root: &Option<Rc<RefCell<TreeNode>>>,
    //     mut current_height: i32,
    //     height: i32,
    // ) -> Option<i32> {
    //     if root.is_none() {
    //         return Some(0);
    //     }
    //     let root_borrow = root.as_ref()?.borrow();
    //     let left = &root_borrow.left;
    //     let right = &root_borrow.right;
    //     if left.is_none() && right.is_none() && current_height == height {
    //         return Some(1);
    //     }
    //     current_height = current_height + 1;
    //     return Some(
    //         Self::get_leaf(left, current_height, height)?
    //             + Self::get_leaf(right, current_height, height)?,
    //     );
    // }
    // fn get_height(root: &Option<Rc<RefCell<TreeNode>>>) -> Option<i32> {
    //     if root.is_none() {
    //         return Some(0);
    //     }
    //     let root_borrow = root.as_ref()?.borrow();
    //     let left = &root_borrow.left;
    //     let right = &root_borrow.right;
    //     return Some(max(Self::get_height(left)?, Self::get_height(right)?) + 1);
    // }
    // 解法 2 bfs
    pub fn count_nodes(root: Option<Rc<RefCell<TreeNode>>>) -> i32 {
        let mut queue = VecDeque::<Option<Rc<RefCell<TreeNode>>>>::new();
        queue.push_back(root);
        let mut res = 0;
        while queue.len() != 0 {
            let mut node_vec = vec![];
            while queue.len() != 0 {
                let node = queue.pop_front().unwrap();
                if node.is_some() {
                    node_vec.push(Self::create_rc(&node)); // 暂存当前行所有元素
                    res += 1; // 保存结果
                }
            }
            for i in &node_vec {
                // 把下一行所有元素入队
                let node_borrow = i.as_ref().unwrap().borrow();
                let left = Self::create_rc(&node_borrow.left);
                let right = Self::create_rc(&node_borrow.right);
                if left.is_some() {
                    queue.push_back(left)
                }
                if right.is_some() {
                    queue.push_back(right)
                }
            }
        }
        res
    }
    fn create_rc(root: &Option<Rc<RefCell<TreeNode>>>) -> Option<Rc<RefCell<TreeNode>>> {
        if root.is_none() {
            return None;
        }
        return Some(Rc::clone(root.as_ref().unwrap()));
    }
}
// @lc code=end