375. Count Nodes Equal to Average of Subtree

A node's subtree average can only be known once its children have reported their own subtree totals, so do the work on the way back up: each call returns the (sum, size) of its subtree.

Idea

Every node needs two facts about its own subtree before it can decide anything: the sum of all values in that subtree and the number of nodes in it. Neither is known while descending — they are only complete once both children have finished. So this is a classic postorder problem: recurse to the children first, then combine on the way back up.

We make dfs(node) return a pair (sum, count) for the subtree rooted at node. At each node we:

1. Get (leftSum, leftCount) and (rightSum, rightCount) from the children.
2. Form total = node.val + leftSum + rightSum and size = leftCount + rightCount + 1.
3. Compare node.val with the floor average total // size; if they match, this node counts.
4. Return (total, size) so the parent can fold it in.

A missing (null) child contributes (0, 0), which is why leaves naturally get total = node.val, size = 1, and always match themselves.

Worked example

Tree [4,8,5,0,1,null,6]:

        4
      /   \
     8     5
    / \     \
   0   1     6

Unwinding from the leaves:

• Node 0: subtree (0, 1), avg 0//1 = 0 → matches 0 ✓
• Node 1: subtree (1, 1), avg 1 → matches 1 ✓
• Node 6: subtree (6, 1), avg 6 → matches 6 ✓
• Node 8: subtree (8+0+1, 3) = (9, 3), avg 9//3 = 3 → value 8 ≠ 3 ✗
• Node 5: subtree (5+6, 2) = (11, 2), avg 11//2 = 5 → matches 5 ✓
• Node 4 (root): subtree (4+9+11, 6) = (24, 6), avg 24//6 = 4 → matches 4 ✓

Five nodes match, so the answer is 5.

Complexity

• Time: O(n) — each node is visited exactly once.
• Space: O(h) — the recursion stack, where h is the height of the tree (O(n) in the worst case of a skewed tree).

class Solution {
    private int count = 0;

    public int averageOfSubtree(TreeNode root) {
        dfs(root);
        return count;
    }

    // Returns {sum, size} for the subtree rooted at node.
    private int[] dfs(TreeNode node) {
        if (node == null) return new int[]{0, 0};
        int[] left = dfs(node.left);
        int[] right = dfs(node.right);
        int total = node.val + left[0] + right[0];
        int size = left[1] + right[1] + 1;
        if (node.val == total / size) count++;
        return new int[]{total, size};
    }
}