Count Nodes Equal to Average of Subtree

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;
    }
};

The provided C++ solution addresses the problem of determining how many nodes in a binary tree are equal to the average of values in their respective subtrees. This implementation defines a Solution class that contains two primary methods: processSubtree and computeAverageOfSubtree, along with an integer attribute totalCount for storing the result.

• Understanding processSubtree method:

  • This function is called recursively and deals with the traversal of each node in the tree. It computes the sum of values and the total number of nodes for each subtree originating from a given node.
  • For each node visited, processSubtree computes the sum of its own value and the values from its left and right subtrees. Similarly, it computes the total count of nodes by counting the node itself and the counts from its left and right subtrees.
  • If the node’s value is equal to the integer division of the sum of its subtree's values by the total number of nodes in its subtree, it increments the totalCount.
  • The function returns a pair consisting of the sum of subtree values and the total number of nodes in the subtree.

• Understanding computeAverageOfSubtree method:

  • This method is designed to be the primary public interface that triggers the computation process. It initializes totalCount to zero then calls processSubtree for the root node and finally returns the value of totalCount.

Executing this solution involves creating an instance of Solution class and invoking the computeAverageOfSubtree with the root node of the binary tree as the argument. This will return the number of nodes whose values are equal to the average of their respective subtrees.

The approach makes use of postorder traversal of the binary tree to calculate sums and counts bottom-up, ensuring that all necessary values for subtree calculations are available when they are needed.

class Solution {
    int totalCount = 0;
    
    Pair<Integer, Integer> calculateSubtree(TreeNode rootNode) {
        if (rootNode == null) {
            return new Pair(0, 0);
        }
        
        Pair<Integer, Integer> leftSubtree = calculateSubtree(rootNode.left);
        Pair<Integer, Integer> rightSubtree = calculateSubtree(rootNode.right);
        
        int subtreeSum = leftSubtree.getKey() + rightSubtree.getKey() + rootNode.val;
        int nodesCount = leftSubtree.getValue() + rightSubtree.getValue() + 1;

        if (rootNode.val == subtreeSum / nodesCount) {
            totalCount++;
        }
        
        return new Pair(subtreeSum, nodesCount);
    }
    
    public int findAverageOfSubtree(TreeNode rootNode) {
        calculateSubtree(rootNode);
        return totalCount;
    }
}

The code snippet in Java is designed to solve the problem of counting nodes in a binary tree that have values equal to the average of their respective subtrees. Here's an overview of how the solution operates:

• The Solution class has two primary components:

  • A field totalCount initialized to zero. This field keeps track of the number of nodes meeting the condition.
  • Two functions: calculateSubtree and findAverageOfSubtree.

• The function calculateSubtree is a recursive method that:

  • Checks if the current node (rootNode) is null. If true, it returns a pair (0,0) indicating zero sum and zero nodes.
  • Recursively calculates the sum of values and the count of nodes for both the left and right subtrees.
  • Computes the total sum and total count for the subtree rooted at the current node by adding the values from both left and right subtrees along with the current node's value.
  • Checks if the value of the current node equals the integer division of the subtree sum by the number of nodes (which gives the average). If they are equal, increments the totalCount.
  • Returns a pair containing the total sum and the count of nodes for the subtree.

• findAverageOfSubtree is the function that invokes calculateSubtree, passing the root of the tree as an argument, and returns totalCount, thus providing the final count of nodes that equal the average of their respective subtrees.

This implementation effectively traverses each node of the tree once, resulting in a time complexity proportional to the number of nodes, O(n), which is efficient for this task. This method leverages the strategy of returning multiple values (sum and count) using the Pair class, which simplifies handling combined recursive computations in the tree structure.