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

2265. Count Nodes Equal to Average of Subtree

Clarification Questions

• No, it's clear from problem description.

Test Cases

Normal Cases

Input: 
Output: 

Edge / Corner Cases

Input: 
Output: 

private var result = 0

fun averageOfSubtree(root: TreeNode?): Int {
    if (root == null) return 0
    dfs(root)
    return result
}

private fun dfs(root: TreeNode): IntArray {
    var sum = root.`val`
    var count = 1
    if (root.left != null) {
        val left = dfs(root.left)
        sum += left[0]
        count += left[1]
    }

    if (root.right != null) {
        val right = dfs(root.right)
        sum += right[0]
        count += right[1]
    }

    if (root.left == null && root.right == null) {
        result++
        return intArrayOf(root.`val`, 1)
    }

    val average = sum / count
    if (average == root.`val`) result++
    return intArrayOf(sum, count)
}

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