Question

2265. Count Nodes Equal to Average of Subtree

中文文档

Description

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

Solutions

Solution 1

# 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
}