easy Count Nodes Equal to Average of Subtree

Problem Statement

Given the root of a binary tree, return the integer value representing the number of nodes where the node's value is equal to the average of the values in its subtree.

A subtree of the root node contains the root node and all of its descendants.

The average value of n elements is totalSum / n and rounded down to the nearest integer.

Examples

Example 1:

• Input: [3,1,4,null,2]

• Expected Output: 3
• Justification: The node with value 3 has an average of (3+1+4+2)/4 = 2.5 (round down to 2), which doesn't match its value. The node with value 1 has an average of (1+2)/2 = 1.5 (round down to 1), matching its value. The node with value 4 has no children, so its average is 4, matching its value. The node with value 2 is a leaf, so its average is 2, matching its value. Therefore, there are 3 nodes whose value matches the average of their subtree.

Example 2:

• Input: [5,4,8,3,null,6,7]

• Expected Output: 4
• Justification: The root node 5 has an average of (5+4+8+3+6+7)/6 = 5.5 (round down to 5), which matches its value. All leaves with values 3, 6, and 7 match their averages (as they have no children). So, there are a total 4 nodes that matches the given criteria.

Example 3:

• Input: [10,20,30,40,50,60,70]

• Expected Output: 4
• Justification: The leaf nodes (40, 50, 60, 70) have their values equal to their averages as they do not have any children. The other nodes' values do not match the average of their subtrees. Therefore, there are 4 nodes whose values match the average of their subtree.

Solution

To solve this problem, we will employ a depth-first search (DFS) algorithm that traverses the entire binary tree to calculate the sum and count of nodes in each subtree. This approach allows us to compute the average value of a subtree at each node by dividing the total sum of the subtree by the total count of nodes in that subtree. The key reason this approach is effective is because it allows for the simultaneous calculation of sum and count at each node in a single traversal, minimizing the computational complexity.

Our algorithm will recursively visit each node, calculate the sum and count of its subtree, and then determine if the node's value matches the calculated average. By traversing the tree in a depth-first manner, we ensure that all descendants of a node are processed before the node itself, which is crucial for accurate average calculation. This method is chosen for its efficiency in handling tree structures and its ability to solve the problem with minimal overhead.

Step-by-Step Algorithm

1. Start at the Root: Begin the analysis with the root node of the binary tree.

2. Recursive Analysis: For each node, recursively analyze both its left and right subtrees to compute the sum of values and the total count of nodes in those subtrees.

3. Handle Base Case: If the current node is null, indicate that there's no contribution to the sum or count from this subtree by returning a result representing zero sum and zero nodes.

4. Compute Subtree Results: For a non-null node, combine the results of the left and right subtree analyses with the current node's value to calculate the total sum and total count of nodes for the subtree rooted at this node.

5. Determine Node Match: Check if the node's value matches the average of the sum of its subtree values. The average is computed as the total sum divided by the total count of nodes, rounded down to the nearest integer. If the node's value matches this average, increment a global count of such matching nodes.

6. Return Subtree Results: After processing a node, return the computed sum and count of nodes for the subtree rooted at this node to its parent.

7. Complete the Analysis: Once all nodes have been analyzed, return the total count of nodes whose values match the average of their respective subtrees.

Algorithm Walkthrough

• The input tree structure:

      5
     / \
    4   8
   /   / \
  3   6   7
  

1. Start at the Root (Value 5): The root's value is 5, and we begin the analysis here, aiming to explore the entire tree.

2. Left Subtree (Value 4):

   • Explore left child (value 4), which itself has a left child (value 3).
   • Node 3 is a leaf, so its average matches its value (3/1 = 3). This node matches the criteria.
   • The subtree rooted at node 4 has a sum of 7 (4 from itself and 3 from its left child) and a total count of 2. The average (7/2 = 3) does not match node 4's value, so it is not counted.

3. Right Subtree (Value 8):

   • Explore right child (value 8), which has two children: 6 (left) and 7 (right).
   • Both children are leaves, so their averages match their values. Nodes 6 and 7 match the criteria, increasing the count.
   • The subtree rooted at node 8 has a sum of 21 (8+6+7) and a count of 3. The average does not match node 8's value, so it is not counted.

4. Back to the Root:

   • After analyzing both subtrees, the root node aggregates these results: the total sum is 33 (5+7+21 from itself and both subtrees) and the total count of nodes is 6.
   • The average for the root's subtree is 33 divided by 6, which is 5.6, rounded down to 6. Since the root's value is 5, it matches the average, and thus, increment the count to 4.

5. Conclusion:

   • There are 4 nodes whose values match the average of their subtree.

Code

// TreeNode definition
// public class TreeNode {
//     int val;
//     TreeNode left;
//     TreeNode right;
//     TreeNode(int x) { val = x; }
// }

class Solution {

  // Method to initiate the subtree analysis
  public int countNodes(TreeNode root) {
    int[] count = new int[1]; // Tracks the number of matching nodes
    subtreeAverage(root, count);
    return count[0];
  }

  // Helper method to calculate the sum and count of nodes in the subtree
  private int[] subtreeAverage(TreeNode node, int[] count) {
    if (node == null) return new int[] { 0, 0 };

    int[] left = subtreeAverage(node.left, count);
    int[] right = subtreeAverage(node.right, count);

    int sum = left[0] + right[0] + node.val;
    int totalCount = left[1] + right[1] + 1;

    if (sum / totalCount == node.val) count[0]++; // Check if node value matches the average

    return new int[] { sum, totalCount };
  }

  // Main method to test the examples
  public static void main(String[] args) {
    Solution sol = new Solution();

    // Example 1
    TreeNode root1 = new TreeNode(3);
    root1.left = new TreeNode(1);
    root1.right = new TreeNode(4);
    root1.left.right = new TreeNode(2);
    System.out.println(sol.countNodes(root1)); // Expected output: 3

    // Example 2
    TreeNode root2 = new TreeNode(5);
    root2.left = new TreeNode(4);
    root2.right = new TreeNode(8);
    root2.left.left = new TreeNode(3);
    root2.right.left = new TreeNode(6);
    root2.right.right = new TreeNode(7);
    System.out.println(sol.countNodes(root2)); // Expected output: 4

    // Example 3
    TreeNode root3 = new TreeNode(10);
    root3.left = new TreeNode(20);
    root3.right = new TreeNode(30);
    root3.left.left = new TreeNode(40);
    root3.left.right = new TreeNode(50);
    root3.right.left = new TreeNode(60);
    root3.right.right = new TreeNode(70);
    System.out.println(sol.countNodes(root3)); // Expected output: 4
  }
}

Complexity Analysis

Time Complexity

• : The algorithm involves a recursive traversal of the binary tree, where N is the number of nodes in the tree. Every node is visited exactly once during the traversal, and for each node, we perform a constant amount of work. Therefore, the time complexity is linear in the number of nodes in the tree.

Space Complexity

• : The space complexity is determined by the height of the binary tree H, due to the recursive call stack. In the worst case (for a skewed tree), the height of the tree can be N (the number of nodes), resulting in a space complexity of . For a balanced tree, the height H would be log(N), resulting in a space complexity of .