【二叉树】统计值等于子树平均值的节点数-CSDN博客

本文链接：https://blog.csdn.net/xjh093/article/details/126173674

0x00 题目

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

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

0x01 思路

要获取某个子树的个数以及所有节点值的和
就需要使用后序遍历方式
总和 = 当前节点值 + 左子树的和 + 右子树的和
个数 = 1 + 左子树的个数 + 右子树的个数

0x02 解法

语言：Swift

树节点：TreeNode

public class TreeNode {
    public var val: Int
    public var left: TreeNode?
    public var right: TreeNode?
    public init() { self.val = 0; self.left = nil; self.right = nil; }
    public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
    public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
        self.val = val
        self.left = left
        self.right = right
    }
}

解法：

func averageOfSubtree(_ root: TreeNode?) -> Int {
    guard let root = root else { return 0 }
    var res: Int = 0
    
    func dfs(_ root: TreeNode?) -> (Int, Int) {
        guard let root = root else { return (0,0) }
        
        let left = dfs(root.left)
        let right = dfs(root.right)
        
        // 后序遍历
        // 总和
        let sum = left.0 + right.0 + root.val
        // 个数
        let count = left.1 + right.1 + 1
        // 平均值
        let average = sum / count
        // 相等则统计
        if (average == root.val) { res += 1}
        
        return (sum, count)
    }

    _ = dfs(root)
    return res
}