Question

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

English Version

题目描述

给你一棵二叉树的根节点 root ，找出并返回满足要求的节点数，要求节点的值等于其 子树 中值的 平均值 。

注意：

• n 个元素的平均值可以由 n 个元素 求和 然后再除以 n ，并 向下舍入 到最近的整数。
• root 的 子树 由 root 和它的所有后代组成。

 

示例 1：

输入：root = [4,8,5,0,1,null,6]
输出：5
解释：
对值为 4 的节点：子树的平均值 (4 + 8 + 5 + 0 + 1 + 6) / 6 = 24 / 6 = 4 。
对值为 5 的节点：子树的平均值 (5 + 6) / 2 = 11 / 2 = 5 。
对值为 0 的节点：子树的平均值 0 / 1 = 0 。
对值为 1 的节点：子树的平均值 1 / 1 = 1 。
对值为 6 的节点：子树的平均值 6 / 1 = 6 。

示例 2：

输入：root = [1]
输出：1
解释：对值为 1 的节点：子树的平均值 1 / 1 = 1。

 

提示：

• 树中节点数目在范围 [1, 1000] 内
• 0 <= Node.val <= 1000

解法

方法一

# Definition for a binary tree node.
# class TreeNode:
#   def __init__(self, val=0, left=None, right=None):
#     self.val = val
#     self.left = left
#     self.right = right
class Solution:
  def averageOfSubtree(self, root: Optional[TreeNode]) -> int:
    def dfs(root):
      if root is None:
        return 0, 0
      ls, ln = dfs(root.left)
      rs, rn = dfs(root.right)
      s = ls + rs + root.val
      n = ln + rn + 1
      if s // n == root.val:
        nonlocal ans
        ans += 1
      return s, n

    ans = 0
    dfs(root)
    return ans

/**
 * 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 {
  private int ans;

  public int averageOfSubtree(TreeNode root) {
    ans = 0;
    dfs(root);
    return ans;
  }

  private int[] dfs(TreeNode root) {
    if (root == null) {
      return new int[] {0, 0};
    }
    int[] l = dfs(root.left);
    int[] r = dfs(root.right);
    int s = l[0] + r[0] + root.val;
    int n = l[1] + r[1] + 1;
    if (s / n == root.val) {
      ++ans;
    }
    return new int[] {s, n};
  }
}

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *   int val;
 *   TreeNode *left;
 *   TreeNode *right;
 *   TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *   TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *   TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
  int ans;
  int averageOfSubtree(TreeNode* root) {
    ans = 0;
    dfs(root);
    return ans;
  }

  vector<int> dfs(TreeNode* root) {
    if (!root) return {0, 0};
    auto l = dfs(root->left);
    auto r = dfs(root->right);
    int s = l[0] + r[0] + root->val;
    int n = l[1] + r[1] + 1;
    if (s / n == root->val) ++ans;
    return {s, n};
  }
};

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *   Val int
 *   Left *TreeNode
 *   Right *TreeNode
 * }
 */
func averageOfSubtree(root *TreeNode) int {
  ans := 0
  var dfs func(*TreeNode) (int, int)
  dfs = func(root *TreeNode) (int, int) {
    if root == nil {
      return 0, 0
    }
    ls, ln := dfs(root.Left)
    rs, rn := dfs(root.Right)
    s := ls + rs + root.Val
    n := ln + rn + 1
    if s/n == root.Val {
      ans++
    }
    return s, n
  }
  dfs(root)
  return ans
}