2265.count-nodes-equal-to-average-of-subtree.md

2265. Count Nodes Equal to Average of Subtree

Bottom-up DFS

We use a bottom-up DFS to traverse the tree. For each node, we calculate the sum of the subtree and the number of nodes in the subtree. We then check if the average of the subtree is equal to the value of the node. If it is, we increment the answer.

private var ans = 0
fun averageOfSubtree(root: TreeNode?): Int {
    dfs(root)
    return ans
}

// Return sum, number of nodes
private fun dfs(root: TreeNode?): Pair<Int, Int> {
    if (root == null) return 0 to 0
    val (leftSum, leftNodes) = dfs(root.left)
    val (rightSum, rightNodes) = dfs(root.right)

    val totalSum = root.`val` + leftSum + rightSum
    val nodes = 1 + leftNodes + rightNodes
    if (totalSum / nodes == root.`val`) ans++

    return totalSum to nodes
}

• Time Complexity: O(n).
• Space Complexity: O(n).