Count Nodes Equal to Average of Subtree

Med

#2063Time: O(N^2) in the worst case. The main `traverse` function visits N nodes. For each node, `getSubtreeInfo` is called, which traverses that node's entire subtree. In a skewed tree (like a linked list), this leads to a complexity of N + (N-1) + ... + 1, which is O(N^2). For a perfectly balanced tree, the complexity is O(N log N).Space: O(N) in the worst case. The space is determined by the maximum depth of the recursion stack. For a skewed tree, the height is N, leading to O(N) space.2 companies

Algorithms

Depth-First Search

Data structures

Tree Binary Tree

Companies

Snowflake INDmoney

Prompt

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

Note:

The average of n elements is the sum of the n elements divided by n and rounded down to the nearest integer.
A subtree of root is a tree consisting of root and all of its descendants.

Example 1:

Input: root = [4,8,5,0,1,null,6]
Output: 5
Explanation: 
For the node with value 4: The average of its subtree is (4 + 8 + 5 + 0 + 1 + 6) / 6 = 24 / 6 = 4.
For the node with value 5: The average of its subtree is (5 + 6) / 2 = 11 / 2 = 5.
For the node with value 0: The average of its subtree is 0 / 1 = 0.
For the node with value 1: The average of its subtree is 1 / 1 = 1.
For the node with value 6: The average of its subtree is 6 / 1 = 6.

Example 2:

Input: root = [1]
Output: 1
Explanation: For the node with value 1: The average of its subtree is 1 / 1 = 1.

Constraints:

The number of nodes in the tree is in the range [1, 1000].
0 <= Node.val <= 1000

Approaches

2 approaches with complexity analysis and trade-offs.

This approach iterates through every node in the tree. For each node, it performs a separate, full traversal of that node's subtree to calculate the sum of values and the number of nodes. It then computes the average and checks if it matches the node's value. This method is straightforward but computationally expensive.

Algorithm

Initialize a global counter result to 0.
Create a main traversal function, traverse(node), that visits every node in the tree (e.g., using pre-order traversal).
For each node visited by traverse:
- Call a separate helper function, getSubtreeInfo(node), to compute the properties of its subtree.
- The getSubtreeInfo function recursively traverses the subtree starting from node to calculate the total sum of node values and the total count of nodes.
- It returns these two values, for example, in an array [sum, count].
- Back in the traverse function, calculate the average: average = sum / count.
- If node.val is equal to average, increment the global result counter.
After the traverse function finishes, return the final result.

Walkthrough

The core idea is to decouple the process of visiting each node from the process of calculating its subtree's properties. We use a primary traversal to iterate through all potential nodes that could satisfy the condition. For each of these nodes, we initiate a secondary, independent traversal on its subtree to gather the necessary data (sum and count) for the average calculation.

1/**2 * Definition for a binary tree node.3 * public class TreeNode {4 *     int val;5 *     TreeNode left;6 *     TreeNode right;7 *     TreeNode(int val) { this.val = val; }8 *     TreeNode(int val, TreeNode left, TreeNode right) {9 *         this.val = val;10 *         this.left = left;11 *         this.right = right;12 *     }13 * }14 */15class Solution {16    int result = 0;17 18    public int averageOfSubtree(TreeNode root) {19        // Main traversal to visit every node20        traverse(root);21        return result;22    }23 24    // This function iterates through each node of the tree25    private void traverse(TreeNode node) {26        if (node == null) {27            return;28        }29 30        // For the current node, get its subtree's sum and count31        int[] subtreeInfo = getSubtreeInfo(node);32        int sum = subtreeInfo[0];33        int count = subtreeInfo[1];34 35        // Check if the node's value equals the average of its subtree36        if (count > 0 && (sum / count) == node.val) {37            result++;38        }39 40        // Continue to the next nodes41        traverse(node.left);42        traverse(node.right);43    }44 45    // This helper function calculates the sum and count for a given subtree46    private int[] getSubtreeInfo(TreeNode subRoot) {47        if (subRoot == null) {48            return new int[]{0, 0}; // {sum, count}49        }50 51        int[] leftInfo = getSubtreeInfo(subRoot.left);52        int[] rightInfo = getSubtreeInfo(subRoot.right);53 54        int totalSum = subRoot.val + leftInfo[0] + rightInfo[0];55        int nodeCount = 1 + leftInfo[1] + rightInfo[1];56 57        return new int[]{totalSum, nodeCount};58    }59}

Complexity

Time

O(N^2) in the worst case. The main `traverse` function visits N nodes. For each node, `getSubtreeInfo` is called, which traverses that node's entire subtree. In a skewed tree (like a linked list), this leads to a complexity of N + (N-1) + ... + 1, which is O(N^2). For a perfectly balanced tree, the complexity is O(N log N).

Space

O(N) in the worst case. The space is determined by the maximum depth of the recursion stack. For a skewed tree, the height is N, leading to O(N) space.

Trade-offs

Pros

Conceptually simple and easy to follow.
Separates the logic of iterating through nodes from the logic of calculating subtree properties.

Cons

Highly inefficient due to massive redundant computations. The sum and count for a subtree are recalculated multiple times as the main traversal moves up the tree.

A more efficient approach uses a single post-order traversal. By processing children before their parent (a bottom-up approach), we can calculate the sum and count for a subtree and pass this information up to the parent. This avoids redundant calculations, as the information for each subtree is computed only once, leading to an optimal linear time solution.

Algorithm

Initialize a global or member variable resultCount to 0.
Create a recursive helper function, postOrder(node), that performs a post-order traversal.
This function will return an array or pair [subtreeSum, subtreeCount] for the subtree rooted at node.
Base Case: If node is null, return [0, 0] as an empty subtree has zero sum and zero count.
Recursive Step:
- Recursively call postOrder on the left child to get [leftSum, leftCount].
- Recursively call postOrder on the right child to get [rightSum, rightCount].
- Now, process the current node (this is the 'post-order' part).
- Calculate the sum and count for its subtree: currentSum = node.val + leftSum + rightSum and currentCount = 1 + leftCount + rightCount.
- Calculate the average: average = currentSum / currentCount.
- If node.val == average, increment resultCount.
- Return [currentSum, currentCount] to the calling function (the parent node).
The main function will call postOrder(root) and then return resultCount.

Walkthrough

The problem has optimal substructure, meaning the solution for a node depends on the solutions for its children. This is a perfect fit for a post-order traversal, which naturally computes results for subproblems (subtrees) before using them to solve the larger problem (the parent's subtree). We design a recursive function that, for any given node, returns the sum and count of its subtree. While doing so, it also performs the required average check for that node.

1/**2 * Definition for a binary tree node.3 * public class TreeNode {4 *     int val;5 *     TreeNode left;6 *     TreeNode right;7 *     TreeNode(int val) { this.val = val; }8 *     TreeNode(int val, TreeNode left, TreeNode right) {9 *         this.val = val;10 *         this.left = left;11 *         this.right = right;12 *     }13 * }14 */15class Solution {16    int resultCount = 0;17 18    public int averageOfSubtree(TreeNode root) {19        postOrder(root);20        return resultCount;21    }22 23    // Returns a pair [sum, count] for the subtree rooted at 'node'24    private int[] postOrder(TreeNode node) {25        if (node == null) {26            return new int[]{0, 0}; // {sum, count}27        }28 29        // 1. Traverse left and right subtrees first (post-order)30        int[] leftSubtree = postOrder(node.left);31        int[] rightSubtree = postOrder(node.right);32 33        // 2. Process the current node using results from children34        int currentSum = node.val + leftSubtree[0] + rightSubtree[0];35        int currentCount = 1 + leftSubtree[1] + rightSubtree[1];36 37        // Calculate average and check the condition38        if (currentSum / currentCount == node.val) {39            resultCount++;40        }41 42        // 3. Return sum and count for the parent's calculation43        return new int[]{currentSum, currentCount};44    }45}

Complexity

Time

O(N), where N is the number of nodes in the tree. Each node is visited and processed exactly once during the single post-order traversal.

Space

O(H), where H is the height of the tree. This space is used by the recursion stack. In the worst case of a skewed tree, H can be N, leading to O(N) space. For a balanced tree, it's O(log N).

Trade-offs

Pros

Optimal time complexity of O(N).
Very efficient as it computes the sum and count for each subtree only once.

Cons

The logic can be slightly less intuitive than the brute-force approach, as the recursive function has to manage and return multiple pieces of information (sum and count).

Solutions

Solution

1/** * Definition for a binary tree node. * public class TreeNode { * int val; * TreeNode left; * TreeNode right; * TreeNode() {} * TreeNode(int val) { this.val = val; } * TreeNode(int val, TreeNode left, TreeNode right) { * this.val = val; * this.left = left; * this.right = right; * } * } */ class Solution { private int ans ; public int averageOfSubtree ( TreeNode root ) { ans = 0 ; dfs ( root ); return ans ; } private int [] dfs ( TreeNode root ) { if ( root == null ) { return new int [] { 0 , 0 }; } int [] l = dfs ( root . left ); int [] r = dfs ( root . right ); int s = l [ 0 ] + r [ 0 ] + root . val ; int n = l [ 1 ] + r [ 1 ] + 1 ; if ( s / n == root . val ) { ++ ans ; } return new int [] { s , n }; } }

1/** * Definition for a binary tree node. * struct TreeNode { * int val; * TreeNode *left; * TreeNode *right; * TreeNode() : val(0), left(nullptr), right(nullptr) {} * TreeNode(int x) : val(x), left(nullptr), right(nullptr) {} * TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {} * }; */ class Solution { public: int ans ; int averageOfSubtree ( TreeNode * root ) { ans = 0 ; dfs ( root ); return ans ; } vector < int > dfs ( TreeNode * root ) { if ( ! root ) return { 0 , 0 }; auto l = dfs ( root -> left ); auto r = dfs ( root -> right ); int s = l [ 0 ] + r [ 0 ] + root -> val ; int n = l [ 1 ] + r [ 1 ] + 1 ; if ( s / n == root -> val ) ++ ans ; return { s , n }; } };

1# Definition for a binary tree node. # class TreeNode: # def __init__(self, val=0, left=None, right=None): # self.val = val # self.left = left # self.right = right class Solution : def averageOfSubtree ( self , root : Optional [ TreeNode ]) -> int : def dfs ( root ): if root is None : return 0 , 0 ls , ln = dfs ( root . left ) rs , rn = dfs ( root . right ) s = ls + rs + root . val n = ln + rn + 1 if s // n == root . val : nonlocal ans ans += 1 return s , n ans = 0 dfs ( root ) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem