Description
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
nelements is the sum of thenelements divided bynand rounded down to the nearest integer. - A subtree of
rootis a tree consisting ofrootand all of its descendants.
Solution
No way this is a medium problem.
This problem could be easily solved with a DFS. The recursive function returns the number of nodes in this subtree start with a node root and the sum of the nodes in the subtree. If the node root has the value equivalent to the average of the nodes in this subtree, add one to result.
Implementation
In my implementation I returned three numbers in the recursive function, which is the number of nodes that satisfies the requirement, the number of nodes and the sum of the values of nodes. We could make the result (# of nodes satisfies the requirement) as a member variable of the Solution class to reduce the data required to return. (But I’m too lazy to spend any more time on this simple problem)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode() : val(0), left(nullptr), right(nullptr) {}
* TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
* TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
* };
*/
class Solution {
public:
vector<int> runner(TreeNode* root) {
vector<int> info(3); // number of qualified nodes, number of nodes, sum of the values of nodes
if (!root) return info;
vector<int> left = runner(root->left);
vector<int> right = runner(root->right);
info[0] = left[0] + right[0];
info[1] = left[1] + right[1] + 1;
info[2] = left[2] + right[2] + root->val;
int avg = info[2] / info[1];
if (root->val == avg) info[0]++;
return info;
}
int averageOfSubtree(TreeNode* root) {
vector<int> result = runner(root);
return result[0];
}
};