Count Nodes Equal to Average of Subtree

medium tree dfs

Problem

Given the root of a binary tree, count how many nodes have a value equal to the average of all values in that node's subtree (the node itself plus every descendant). The average is computed with integer division — divide the subtree's total sum by the number of nodes in the subtree and drop any fractional part.

Inputtree (level-order) = [4, 8, 5, 0, 1, _, 6]

Output5

Root 4: subtree sum 24, size 6, avg 4 → match. Node 5: sum 11, size 2, avg 5 → match. The three leaves 0, 1, 6 each equal their own average. Node 8: sum 9, size 3, avg 3 → no. So 5 nodes qualify.

tree (level-order, '_' for null)

def average_of_subtree(root):
    count = 0
    def dfs(n):
        nonlocal count
        if not n: return (0, 0)
        ls, lc = dfs(n.left)
        rs, rc = dfs(n.right)
        s = ls + rs + n.val
        c = lc + rc + 1
        if n.val == s // c: count += 1
        return (s, c)
    dfs(root)
    return count

function averageOfSubtree(root) {
  let count = 0;
  function dfs(n) {
    if (!n) return [0, 0];
    const [ls, lc] = dfs(n.left);
    const [rs, rc] = dfs(n.right);
    const s = ls + rs + n.val;
    const c = lc + rc + 1;
    if (n.val === Math.floor(s / c)) count++;
    return [s, c];
  }
  dfs(root);
  return count;
}

class Solution {
    int count;
    public int averageOfSubtree(TreeNode root) {
        count = 0;
        dfs(root);
        return count;
    }
    private int[] dfs(TreeNode n) {
        if (n == null) return new int[]{0, 0};
        int[] L = dfs(n.left);
        int[] R = dfs(n.right);
        int s = L[0] + R[0] + n.val;
        int c = L[1] + R[1] + 1;
        if (n.val == s / c) count++;
        return new int[]{s, c};
    }
}

int count_;
pair<int, int> dfs(TreeNode* n) {
    if (!n) return {0, 0};
    auto L = dfs(n->left);
    auto R = dfs(n->right);
    int s = L.first + R.first + n->val;
    int c = L.second + R.second + 1;
    if (n->val == s / c) count_++;
    return {s, c};
}
int averageOfSubtree(TreeNode* root) {
    count_ = 0;
    dfs(root);
    return count_;
}

Explanation

To know whether a node matches its subtree average, we need two facts about that subtree: the total sum of all values in it and the number of nodes in it. With those two, the average is just sum // count (integer division), and we compare it against the node's own value.

The clean way to gather both at once is a single post-order DFS. We visit children before the parent, and each call returns the pair (sum, count) for the subtree rooted at that node. An empty (null) child returns (0, 0).

At a real node we combine the left pair and the right pair with the node itself: s = leftSum + rightSum + n.val and c = leftCount + rightCount + 1. Now we check n.val == s // c; if it holds, we bump a shared count. Then we hand the freshly computed (s, c) back up to our parent so it can do the same.

Example: tree [4, 8, 5, 0, 1, _, 6]. The leaves 0, 1, and 6 trivially equal their own averages. Node 5 has subtree sum 5 + 6 = 11 over 2 nodes, average 5 — a match. Node 8 has sum 8 + 0 + 1 = 9 over 3 nodes, average 3 — no match. The root 4 has sum 24 over all 6 nodes, average 4 — a match. That is 5 qualifying nodes.

Because each node's (sum, count) is computed exactly once and reused by its parent, the whole tree is processed in a single linear pass.

Time: O(n) Space: O(h)