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
}