Count Nodes Equal to Average of Subtree LeetCode solution in C++

1
class Solution {
2
 public:
3
  pair<int, int> func(TreeNode *root, int &ans) {
4
    if (!root) return {0, 0};
5
    auto p1 = func(root->left, ans);
6
    auto p2 = func(root->right, ans);
7
    int avg = (root->val + p1.first + p2.first) / (p1.second + p2.second + 1);
8
    if (avg == root->val) ans++;
9
    return {root->val + p1.first + p2.first, p1.second + p2.second + 1};
10
  }
11
  int averageOfSubtree(TreeNode *root) {
12
    int ans = 0;
13
    func(root, ans);
14
    return ans;
15
  }
16
};