Medium

题目描述

给定一棵二叉树的根节点 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

解题思路

这道题需要计算每个节点的子树平均值，并统计有多少个节点的值等于其子树平均值。

解题思路

要计算子树的平均值，我们需要两个关键信息：

  1. 子树中所有节点值的总和
  2. 子树中节点的数量

最直观的方法是使用后序遍历（左右根），因为我们需要先计算左右子树的信息，再处理当前节点。

算法步骤

  1. 使用深度优先搜索（DFS）遍历二叉树
  2. 对于每个节点，递归计算其左右子树的总和与节点数量
  3. 当前子树的总和 = 当前节点值 + 左子树总和 + 右子树总和
  4. 当前子树的节点数 = 1 + 左子树节点数 + 右子树节点数
  5. 计算平均值并与当前节点值比较，如果相等则计数器加1
  6. 返回当前子树的总和、节点数量等信息供上层调用

实现方式

我们可以定义一个辅助函数返回一个包含总和、节点数量的元组/结构体，同时在递归过程中更新答案计数器。这样只需要一次遍历就能完成所有计算。

时间复杂度为 O(n)，空间复杂度为 O(h)，其中 h 是树的高度。

代码实现

class Solution {
public:
    int result = 0;
    
    pair<int, int> dfs(TreeNode* root) {
        if (!root) return {0, 0};
        
        auto [leftSum, leftCount] = dfs(root->left);
        auto [rightSum, rightCount] = dfs(root->right);
        
        int totalSum = leftSum + rightSum + root->val;
        int totalCount = leftCount + rightCount + 1;
        
        if (totalSum / totalCount == root->val) {
            result++;
        }
        
        return {totalSum, totalCount};
    }
    
    int averageOfSubtree(TreeNode* root) {
        dfs(root);
        return result;
    }
};

class Solution:
    def averageOfSubtree(self, root: TreeNode) -> int:
        self.result = 0
        
        def dfs(node):
            if not node:
                return 0, 0
            
            left_sum, left_count = dfs(node.left)
            right_sum, right_count = dfs(node.right)
            
            total_sum = left_sum + right_sum + node.val
            total_count = left_count + right_count + 1
            
            if total_sum // total_count == node.val:
                self.result += 1
            
            return total_sum, total_count
        
        dfs(root)
        return self.result

public class Solution {
    private int result = 0;
    
    public int AverageOfSubtree(TreeNode root) {
        DFS(root);
        return result;
    }
    
    private (int sum, int count) DFS(TreeNode node) {
        if (node == null) return (0, 0);
        
        var (leftSum, leftCount) = DFS(node.left);
        var (rightSum, rightCount) = DFS(node.right);
        
        int totalSum = leftSum + rightSum + node.val;
        int totalCount = leftCount + rightCount + 1;
        
        if (totalSum / totalCount == node.val) {
            result++;
        }
        
        return (totalSum, totalCount);
    }
}

var averageOfSubtree = function(root) {
    let count = 0;
    
    function dfs(node) {
        if (!node) return [0, 0];
        
        let [leftSum, leftCount] = dfs(node.left);
        let [rightSum, rightCount] = dfs(node.right);
        
        let totalSum = leftSum + rightSum + node.val;
        let totalCount = leftCount + rightCount + 1;
        
        if (Math.floor(totalSum / totalCount) === node.val) {
            count++;
        }
        
        return [totalSum, totalCount];
    }
    
    dfs(root);
    return count;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型复杂度说明
时间复杂度O(n)需要遍历每个节点一次，其中 n 是树中节点的数量
空间复杂度O(h)递归调用栈的深度为树的高度 h，最坏情况下为 O(n)

相关题目