Count Nodes Equal to Average of Subtree

Problem Statement

The challenge is to determine how many nodes in a given binary tree have values that match the average of all the values in their respective subtrees. A subtree here is defined as a node along with all its descendants. The average is computed by summing all values and dividing by the count, with the result floored to the nearest whole number. The problem relies on efficiently traversing the tree and computing these averages dynamically, ensuring the solution scales well with larger trees.

Examples

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

Approach and Intuition

1. Tree Traversal Mechanism:
   To solve this problem, a depth-first search (DFS) is appropriate as it allows us to access each node and its entire subtree seamlessly. This traversal enables the calculation of each node’s subtree average during the backtracking phase of DFS.

2. Calculate Subtree Sum and Count:

   • During the visit to each node, accumulate both the sum of the subtree values and the count of nodes in the subtree.
   • Each node returns its cumulative sum and count to its parent, which then adds these to its own sum and count.

3. Compute and Compare Tree Averages:

   • Once the sum and count for a node's subtree are known, compute the average by dividing the sum by the count and applying the floor function.
   • Directly compare this average to the node’s value to see if they match.

4. Count Matching Nodes:

   • Maintain a count of nodes where the node value equals the computed subtree average.
   • This counter is updated during the DFS traversal whenever a match is found.

5. Returning the Result:

   • After the DFS traversal completes, the counter will contain the number of nodes where the node value equals the average of its subtree.

This approach ensures that each node is processed once, and the subtree calculations are combined efficiently, leveraging the inherent recursive structure of DFS. The problem constraints allow this DFS approach to run within acceptable limits for time complexity.

Solutions

class Solution {
public:
    int totalCount = 0;
        
    pair<int, int> processSubtree(TreeNode* node) {
        if (node == NULL) {
            return {0, 0};
        }
            
        pair<int, int> leftSubtree = processSubtree(node->left);
        pair<int, int> rightSubtree = processSubtree(node->right);
            
        int sumValues = leftSubtree.first + rightSubtree.first + node->val;
        int totalNodes = leftSubtree.second + rightSubtree.second + 1;
    
        if (node->val == sumValues / totalNodes) {
            totalCount++;
        }
            
        return {sumValues, totalNodes};
    }
        
    int computeAverageOfSubtree(TreeNode* node) {
        processSubtree(node);
        return totalCount;
    }
};