Count Nodes Equal to Average of Subtree - Practice Coding Problems

Count Nodes Equal to Average of Subtree - Problem

Given the root of a binary tree, you need to count how many nodes have a special property: the node's value equals the average of all values in its subtree.

🌳 What is a subtree? For any node, its subtree includes the node itself and all of its descendants (children, grandchildren, etc.).

📊 How is the average calculated? Take the sum of all values in the subtree, divide by the count of nodes, and round down to the nearest integer (floor division).

Example: If a subtree has values [4, 8, 5], the average is (4+8+5)/3 = 17/3 = 5.67... → rounds down to 5

Your task is to return the total count of nodes that satisfy this condition.

Input & Output

example_1.py — Basic Tree

$ Input: root = [4,8,5,0,1,null,6]

› Output: 5

💡 Note: Node 4: subtree [4,8,5,0,1,6] → sum=24, count=6, avg=4 ✓ Node 8: subtree [8,0,1] → sum=9, count=3, avg=3 ✗ Node 5: subtree [5,6] → sum=11, count=2, avg=5 ✓ Node 0: subtree [0] → sum=0, count=1, avg=0 ✓ Node 1: subtree [1] → sum=1, count=1, avg=1 ✓ Node 6: subtree [6] → sum=6, count=1, avg=6 ✓

example_2.py — Single Node

$ Input: root = [1]

› Output: 1

💡 Note: Single node has subtree containing only itself. sum=1, count=1, average=1, so 1==1 is true.

example_3.py — All Different

$ Input: root = [1,2,3]

› Output: 2

💡 Note: Node 1: subtree [1,2,3] → sum=6, count=3, avg=2 ✗ Node 2: subtree [2] → sum=2, count=1, avg=2 ✓ Node 3: subtree [3] → sum=3, count=1, avg=3 ✓

Constraints

The number of nodes in the tree is in the range [1, 1000]
0 ≤ Node.val ≤ 1000
The tree is guaranteed to be a valid binary tree
Average is calculated using integer division (floor division)

Visualization

Tap to expand

Understanding the Visualization

Leaf Employees Report

Each leaf employee reports: 'My team size is 1, total salary is my salary, and I'm valid if my salary equals the average (which is just my salary).'

Middle Managers Aggregate

Each manager combines reports from subordinates: 'My team size = 1 + subordinates' team sizes, total salary = my salary + subordinates' totals.'

Check Manager Validity

Each manager calculates: 'Average = total salary ÷ team size. Am I valid? My salary == average?'

CEO Gets Final Count

The CEO (root) receives the total count of all valid managers in the entire organization.

Key Takeaway

🎯 Key Insight: Post-order traversal allows each node to collect complete information from its subtree before making its own decision, enabling a single-pass solution.

Asked in

a Amazon 45 G Google 38 ⊞ Microsoft 32 f Meta 28 Apple 22

The optimal solution uses a single post-order DFS traversal where each node returns three values: subtree sum, node count, and count of valid nodes. This eliminates redundant calculations and achieves O(n) time complexity compared to the brute force O(n²) approach that recalculates subtree information for each node.

Common Approaches

Approach	Time	Space	Notes
✓ Single DFS Traversal (Optimal)	O(n)	O(h)	Post-order DFS returning sum, count, and result count in one pass
Brute Force (Nested Traversals)	O(n²)	O(h)	For each node, separately calculate subtree sum and count

Single DFS Traversal (Optimal) — Algorithm Steps

Use post-order DFS (left, right, root)
Each recursive call returns (sum, count, validNodes)
For each node, calculate subtree sum and count from children
Check if current node's value equals sum // count
Return updated sum, count, and validNodes count to parent
Base case: null nodes return (0, 0, 0)

Visualization

Tap to expand

Step-by-Step Walkthrough

Process Leaves First

Calculate sum=0, count=1, valid=1 for leaf nodes

Process Internal Nodes

Combine children's results and check current node

Propagate Results Up

Each level passes combined results to parent

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>

struct TreeNode {
    int val;
    struct TreeNode *left;
    struct TreeNode *right;
};

typedef struct {
    int sum;
    int count;
    int validNodes;
} SubtreeResult;

SubtreeResult dfs(struct TreeNode* node) {
    SubtreeResult result = {0, 0, 0};
    
    if (!node) {
        return result;
    }
    
    // Post-order: process children first
    SubtreeResult left = dfs(node->left);
    SubtreeResult right = dfs(node->right);
    
    // Calculate current subtree info
    int currentSum = left.sum + right.sum + node->val;
    int currentCount = left.count + right.count + 1;
    
    // Check if current node is valid
    int average = currentSum / currentCount;
    int currentValid = left.validNodes + right.validNodes;
    if (node->val == average) {
        currentValid++;
    }
    
    result.sum = currentSum;
    result.count = currentCount;
    result.validNodes = currentValid;
    
    return result;
}

int averageOfSubtree(struct TreeNode* root) {
    SubtreeResult result = dfs(root);
    return result.validNodes;
}

int main() {
    printf("Solution ready for testing\n");
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n)

Visit each node exactly once in post-order traversal

✓ Linear Growth

Space Complexity

O(h)

Recursion stack depth equals tree height h, O(log n) for balanced trees, O(n) for skewed trees

✓ Linear Space

42.5K Views

Medium Frequency

~18 min Avg. Time

1.8K Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Single DFS Traversal (Optimal) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler