2265. Count Nodes Equal to Average of Subtree Medium

1/**
2 * [2265] Count Nodes Equal to Average of Subtree
3 *
4 * Given the root of a binary tree, return the number of nodes where the value of the node is equal to the average of the values in its subtree.
5 * Note:
6 * 
7 * The average of n elements is the sum of the n elements divided by n and rounded down to the nearest integer.
8 * A subtree of root is a tree consisting of root and all of its descendants.
9 * 
10 *  
11 * Example 1:
12 * <img src="https://assets.leetcode.com/uploads/2022/03/15/image-20220315203925-1.png" style="width: 300px; height: 212px;" />
13 * Input: root = [4,8,5,0,1,null,6]
14 * Output: 5
15 * Explanation: 
16 * For the node with value 4: The average of its subtree is (4 + 8 + 5 + 0 + 1 + 6) / 6 = 24 / 6 = 4.
17 * For the node with value 5: The average of its subtree is (5 + 6) / 2 = 11 / 2 = 5.
18 * For the node with value 0: The average of its subtree is 0 / 1 = 0.
19 * For the node with value 1: The average of its subtree is 1 / 1 = 1.
20 * For the node with value 6: The average of its subtree is 6 / 1 = 6.
21 * 
22 * Example 2:
23 * <img src="https://assets.leetcode.com/uploads/2022/03/26/image-20220326133920-1.png" style="width: 80px; height: 76px;" />
24 * Input: root = [1]
25 * Output: 1
26 * Explanation: For the node with value 1: The average of its subtree is 1 / 1 = 1.
27 * 
28 *  
29 * Constraints:
30 * 
31 * The number of nodes in the tree is in the range [1, 1000].
32 * 0 <= Node.val <= 1000
33 * 
34 */
35pub struct Solution {}
36#[allow(unused_imports)]
37use crate::util::tree::{TreeNode, to_tree};
38
39// problem: https://leetcode.com/problems/count-nodes-equal-to-average-of-subtree/
40// discuss: https://leetcode.com/problems/count-nodes-equal-to-average-of-subtree/discuss/?currentPage=1&orderBy=most_votes&query=
41
42// submission codes start here
43
44// Definition for a binary tree node.
45// #[derive(Debug, PartialEq, Eq)]
46// pub struct TreeNode {
47//   pub val: i32,
48//   pub left: Option<Rc<RefCell<TreeNode>>>,
49//   pub right: Option<Rc<RefCell<TreeNode>>>,
50// }
51// 
52// impl TreeNode {
53//   #[inline]
54//   pub fn new(val: i32) -> Self {
55//     TreeNode {
56//       val,
57//       left: None,
58//       right: None
59//     }
60//   }
61// }
62use std::rc::Rc;
63use std::cell::RefCell;
64impl Solution {
65    pub fn average_of_subtree(root: Option<Rc<RefCell<TreeNode>>>) -> i32 {
66        let result = Self::inner(root);
67        result.2
68    }
69
70    // (total, count, anser)
71    fn inner(root: Option<Rc<RefCell<TreeNode>>>) -> (i32, i32, i32) {
72        match root {
73            None => {
74                println!("a null");
75                (0, 0, 0)
76            },
77            Some(rc) => {
78                let x = rc.borrow();
79                let left = match x.left.as_ref() {
80                    None => (0, 0, 0),
81                    Some(l) => Self::inner(Some(Rc::clone(l)))
82                };
83                let right = match x.right.as_ref() {
84                    None => (0, 0, 0),
85                    Some(r) => Self::inner(Some(Rc::clone(r)))
86                };
87
88                if left.1 == 0 && right.1 == 0 {
89                    return (x.val, 1, 1);
90                } else {
91                    (left.0 + right.0 + x.val, 
92                        left.1 + right.1 + 1,
93                         if (left.0 + right.0 + x.val) / (left.1 + right.1 + 1) == x.val {1} else {0} + left.2 + right.2)
94                }
95            }
96        }
97    }
98}
99
100// submission codes end
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use crate::tree;
106
107    #[test]
108    fn test_2265_1() {
109        let t = tree!(4,8,5,0,1,null,6);
110        println!("{:#?}", t);
111        assert_eq!(Solution::average_of_subtree(t), 5);
112    }
113}
114