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
Solutions
Solution 1: DFS
We design a function $\textit{dfs}$, which calculates the sum and the number of nodes of the subtree rooted at the current node.
The execution process of the function $\textit{dfs}$ is as follows:
If the current node is null, return $(0, 0)$.
Otherwise, we recursively calculate the sum and the number of nodes of the left and right subtrees, denoted as $(\textit{ls}, \textit{ln})$ and $(\textit{rs}, \textit{rn})$, respectively. Then, the sum $\textit{s}$ and the number of nodes $\textit{n}$ of the subtree rooted at the current node are $\textit{ls} + \textit{rs} + \textit{root.val}$ and $\textit{ln} + \textit{rn} + 1$, respectively. If $\textit{s} / \textit{n} = \textit{root.val}$, it means the current node meets the requirement of the problem, and we increment the answer $\textit{ans}$ by $1$.
Finally, the function $\textit{dfs}$ returns $\textit{s}$ and $\textit{n}$.
We initialize the answer $\textit{ans}$ to $0$, then call the $\textit{dfs}$ function, and finally return the answer $\textit{ans}$.
The time complexity is $O(n)$, and the space complexity is $O(n)$. Here, $n$ represents the number of nodes in the binary tree.
/** * Definition for a binary tree node. * type TreeNode struct { * Val int * Left *TreeNode * Right *TreeNode * } */funcaverageOfSubtree(root*TreeNode)(ansint){vardfsfunc(root*TreeNode)(int,int)dfs=func(root*TreeNode)(int,int){ifroot==nil{return0,0}ls,ln:=dfs(root.Left)rs,rn:=dfs(root.Right)s,n:=ls+rs+root.Val,ln+rn+1ifs/n==root.Val{ans++}returns,n}dfs(root)return}