[LeetCode] 2265. Count Nodes Equal to Average of Subtree

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.

Note:
The average of n elements is the sum of the n elements divided by n and rounded down to the nearest integer.
A subtree of root is a tree consisting of root and all of its descendants.

Example 1:
Input: root = [4,8,5,0,1,null,6]
Output: 5
Explanation:
For the node with value 4: The average of its subtree is (4 + 8 + 5 + 0 + 1 + 6) / 6 = 24 / 6 = 4.
For the node with value 5: The average of its subtree is (5 + 6) / 2 = 11 / 2 = 5.
For the node with value 0: The average of its subtree is 0 / 1 = 0.
For the node with value 1: The average of its subtree is 1 / 1 = 1.
For the node with value 6: The average of its subtree is 6 / 1 = 6.

Example 2:
Input: root = [1]
Output: 1
Explanation: For the node with value 1: The average of its subtree is 1 / 1 = 1.

Constraints:
The number of nodes in the tree is in the range [1, 1000].
0 <= Node.val <= 1000

统计值等于子树平均值的节点数。

给你一棵二叉树的根节点 root ，找出并返回满足要求的节点数，要求节点的值等于其 子树 中值的 平均值 。
注意：
n 个元素的平均值可以由 n 个元素 求和 然后再除以 n ，并 向下舍入 到最近的整数。
root 的 子树 由 root 和它的所有后代组成。

思路

思路是后序遍历。因为题目让我们找的节点需要让其子树满足一些条件，这也就意味着对于任何一个节点，他需要从他的子节点获取一些信息，由此我们想到大致的思路是后序遍历。

那么具体一点，子节点需要往父节点传什么信息呢？因为要求的是子树中所有节点 val 的平均值，那么我们需要收集所有子节点的 node.val 和子节点的个数这两个信息。这里我用一个数组保存。其他部分就比较直观了。

复杂度

时间O(n)
空间O(n)

代码

Java实现

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    int count;

    public int averageOfSubtree(TreeNode root) {
        count = 0;
        if (root == null) {
            return count;
        }
        helper(root);
        return count;
    }

    // [sum, count of nodes]
    private int[] helper(TreeNode root) {
        if (root == null) {
            return new int[] { 0, 0 };
        }
        int[] left = helper(root.left);
        int[] right = helper(root.right);
        int sum = left[0] + right[0] + root.val;
        int nodeCount = left[1] + right[1] + 1;
        if (root.val == sum / nodeCount) {
            count++;
        }
        return new int[] { sum, nodeCount };
    }
}