2265. Count Nodes Equal to Average of Subtree ⭐

Medium

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

解題

原本的程式碼又醜又笨，討論區這個解答我最喜歡：

Runtime: 0 ms, faster than 100.00%

Memory Usage: 4.3 MB, less than 92.31%

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func averageOfSubtree(root *TreeNode) int {
    _, _, r := helper(root)
    return r
}

func helper(root *TreeNode) (count int, sum int, result int) {
    if root == nil {
        return 0, 0, 0
    }
    
    cL, sL, rL := helper(root.Left)
    cR, sR, rR := helper(root.Right)
    
    count = cL + cR + 1
    sum = sL + sR + root.Val
    result = rL+rR
    if sum / count == root.Val {
        result++
    }
    
    return
}