LogoKode$word

LeetCode 2265: Count Nodes Equal to Average of Subtree – Java DFS Solution

Understand subtree sum and node counting with DFS, explore a straightforward recursive solution, optimize it using a shared counter, and finally reach an O(n) postorder solution.

Krishna Shrivastava
12 views
LinkedInGithubX
0
0
LeetCodeJavaBinary TreeDFSRecursionPostorder TraversalMedium
LeetCode 2265: Count Nodes Equal to Average of Subtree – Java DFS Solution
Listen to articleAudio version
Ad

Introduction

Binary tree problems often become interesting when the answer for a node depends on everything below that node.

This problem is exactly that.

For every node, the task is to calculate the average of all values in its subtree and check whether that average is equal to the value of the node itself.

At first, this looks like a simple DFS problem. But there is an important optimization opportunity hidden inside it.

To calculate the average of a subtree, two things are required:

  1. The sum of all nodes in the subtree
  2. The number of nodes in the subtree

The first approach calculates these values directly. Then, by looking at the repeated work, the solution can be improved further.

Question Link

LeetCode 2265 – Count Nodes Equal to Average of Subtree

Understanding the Problem

For every node:

Average = Subtree Sum / Number of Nodes

The average is rounded down.

If the calculated average is equal to the current node's value, that node is counted.

Consider this tree:

4
/ \
8 5
/ \ \
0 1 6

For node 5:

Subtree = [5, 6]

Sum = 5 + 6 = 11
Count = 2

Average = 11 / 2
= 5

Since the average is 5, node 5 contributes to the answer.

The same calculation has to be performed for every node.

The Core Observation

For each node, the required information is:

Subtree Sum
+
Subtree Node Count

Once these two values are available, calculating the average is straightforward.

The natural question is:

How can the sum and count of every subtree be calculated efficiently?

A recursive DFS is a natural fit because the information of a subtree can be obtained by processing its children first.

Approach 1: Calculate Subtree Sum and Count Separately

The first approach is the most direct way to think about the problem.

For every node:

  1. Calculate the sum of its entire subtree.
  2. Calculate the number of nodes in that subtree.
  3. Calculate the average.
  4. Compare it with the current node.
  5. Repeat for the left and right children.

A separate count() function can be used to count the nodes.

This approach follows the problem statement very closely, which makes it a good starting point for understanding the solution.

Code

Here is the complete approach, including the separate count() idea:

class Solution {

int co = 0;

// Calculates the sum of all nodes in the subtree
public int sumdfs(TreeNode roo) {

if (roo == null) {
return 0;
}

// Leaf node contributes its value to the sum
if (roo.left == null && roo.right == null) {
co++;
return roo.val;
}

int left = sumdfs(roo.left);
int right = sumdfs(roo.right);

// Current node is also part of its subtree
co++;

return roo.val + left + right;
}

// Separate function to count the number of nodes
// in the subtree.
//
// This was the initial idea, but it is commented out
// because the sumdfs() traversal can already count nodes.
//
// public int count(TreeNode roo){
// if(roo == null){
// return 0;
// }
// if(roo.left == null && roo.right == null){
// return 1;
// }
// int left = count(roo.left);
// int right = count(roo.right);
// return 1 + left + right;
// }

int am = 0;

public void solve(TreeNode root) {

if (root == null) {
return;
}

// Reset count for the current subtree
co = 0;

// Calculate the sum of the current subtree
// and count its nodes at the same time.
int tot = sumdfs(root);

int cot = co;

// The problem requires the average to be rounded down.
if (Math.floor(tot / cot) == root.val) {
am++;
}

// Check the left subtree
solve(root.left);

// Check the right subtree
solve(root.right);
}

public int averageOfSubtree(TreeNode root) {

solve(root);

return am;
}
}

How This Approach Works

The interesting part is sumdfs().

Initially, a separate count() function can be used:

int cot = count(root);

At the same time, another traversal calculates the sum.

But there is no need to traverse the same subtree again just to count its nodes.

The sumdfs() function is already visiting every node, so the counter can be increased during the same traversal.

That is why the following idea works:

co++;

Whenever a node is visited, co increases by one.

After sumdfs(root) finishes:

int tot = sumdfs(root);
int cot = co;

Now:

tot = sum of subtree
cot = number of nodes in subtree

and the average can be calculated.

Dry Run

Consider:

4
/ \
8 5
/ \ \
0 1 6

Start from node 4.

Node 0

sum = 0
count = 1
average = 0 / 1 = 0

So node 0 is counted.

Node 1

sum = 1
count = 1
average = 1 / 1 = 1

Node 1 is counted.

Node 8

Its subtree is:

8
/ \
0 1

Therefore:

sum = 8 + 0 + 1
= 9

count = 3

average = 9 / 3
= 3

Since:

8 != 3

node 8 is not counted.

Node 6

sum = 6
count = 1

average = 6

Node 6 is counted.

Node 5

Its subtree is:

5
\
6

So:

sum = 5 + 6
= 11

count = 2

average = 11 / 2
= 5

Node 5 is counted.

Node 4

The complete subtree is:

4
/ \
8 5
/ \ \
0 1 6

Therefore:

sum = 4 + 8 + 5 + 0 + 1 + 6
= 24

count = 6

average = 24 / 6
= 4

Node 4 is counted.

The final answer is:

5

An Important Complexity Observation

The co variable removes the need for a separate count() traversal.

However, there is still a bigger issue.

Look at what happens when solve() processes the root.

It calls:

sumdfs(root);

which visits the entire tree.

Then solve() moves to the left child and calls:

sumdfs(root.left);

which visits that subtree again.

Then the same thing happens for the right subtree.

So the same nodes can be visited many times.

For example, in a skewed tree:

1
\
2
\
3
\
4
\
5

The first call processes:

1 → 2 → 3 → 4 → 5

The second processes:

2 → 3 → 4 → 5

The third:

3 → 4 → 5

and so on.

Therefore, although the co optimization is cleaner than using a separate count() traversal, the overall worst-case complexity can still reach:

O(n²)

This observation leads directly to the optimal solution.

Approach 2: Calculate Everything During One Postorder Traversal

Instead of recalculating every subtree, the better idea is:

Calculate the subtree information once and return it to the parent.

For every node, the DFS can return two values:

sum
count

Suppose a node has:

left subtree → (leftSum, leftCount)
right subtree → (rightSum, rightCount)

Then the current node can calculate:

sum = root.val + leftSum + rightSum

and:

count = 1 + leftCount + rightCount

Now the average is immediately available.

After checking the current node, the pair:

(sum, count)

is returned to its parent.

This means every node is processed exactly once.

Why Postorder Traversal?

The traversal order is:

Left → Right → Root

This is called postorder traversal.

It fits perfectly because the parent needs information from its children before it can calculate its own subtree information.

For example:

4
/ \
8 5

The algorithm first calculates information for:

8
5

and then combines them at:

4

This is a very common pattern in tree problems.

Optimal Java Solution

A small Pair class can be used to return both the subtree sum and node count together.

class Solution {

int ans = 0;

class Pair {
int sum;
int count;

Pair(int sum, int count) {
this.sum = sum;
this.count = count;
}
}

public Pair dfs(TreeNode root) {

// Empty subtree
if (root == null) {
return new Pair(0, 0);
}

// Get sum and count from left subtree
Pair left = dfs(root.left);

// Get sum and count from right subtree
Pair right = dfs(root.right);

// Calculate information for current subtree
int sum = root.val + left.sum + right.sum;
int count = 1 + left.count + right.count;

// Calculate subtree average
int average = sum / count;

// Check if current node equals its subtree average
if (average == root.val) {
ans++;
}

// Return current subtree information to parent
return new Pair(sum, count);
}

public int averageOfSubtree(TreeNode root) {

dfs(root);

return ans;
}
}

Why This Solution Is Optimal

Consider node 4.

Instead of recalculating its subtree from scratch, the algorithm receives:

From 8:
sum = 9
count = 3

and:

From 5:
sum = 11
count = 2

Then:

sum = 4 + 9 + 11
= 24

and:

count = 1 + 3 + 2
= 6

Therefore:

average = 24 / 6
= 4

The information is calculated once and passed upward.

No subtree needs to be recalculated.

Complexity Analysis

Approach 1

The sumdfs() traversal can be repeated for every node.

Time Complexity

O(n²)

in the worst case.

Space Complexity

The recursion stack depends on the height of the tree:

O(h)

Approach 2: One Postorder DFS

Every node is visited exactly once.

At each node, only constant-time operations are performed.

Time Complexity

O(n)

Space Complexity

O(h)

where h is the height of the tree.

For a balanced tree:

O(log n)

For a skewed tree:

O(n)

A Small Java Improvement

There is also a small simplification in the original code.

The following expression:

Math.floor(tot / cot)

is unnecessary here.

Because tot and cot are integers:

tot / cot

already performs integer division.

For example:

11 / 2 = 5

which is exactly the required rounded-down result because all node values are non-negative.

Therefore, the comparison can simply be:

if (tot / cot == root.val)

The optimal solution uses:

int average = sum / count;

which is cleaner and easier to read.

Approach Comparison

ApproachMain IdeaTimeSpace
Separate sum + countCalculate subtree information directlyO(n²) worst caseO(h)
Global co counterCount nodes during sum traversalO(n²) worst caseO(h)
One postorder DFSReturn sum + count to parentO(n)O(h)

The second approach is a useful improvement over the first, but the third approach is the real optimization.

What Can Be Learned From This Problem?

This problem teaches an important binary-tree pattern:

If a parent needs information about its entire subtree, calculate that information from the results returned by its children.

Instead of repeatedly asking:

"What is the sum of this subtree?"
"What is the count of this subtree?"

the children calculate their information once and return it.

The parent simply combines the results.

This pattern appears in many tree problems involving:

  1. Subtree sums
  2. Subtree sizes
  3. Tree height
  4. Diameter
  5. Balanced tree checking
  6. Maximum path calculations
  7. Counting nodes satisfying a condition
  8. Tree dynamic programming

A useful mental template is:

Result dfs(Node root) {

if (root == null) {
return baseResult;
}

Result left = dfs(root.left);
Result right = dfs(root.right);

// Combine left + right + current node

return result;
}

Once this pattern becomes familiar, many seemingly complicated tree problems become much easier to approach.

Interview Tip

If an interviewer asks for the straightforward solution first, it is perfectly reasonable to start with the direct recursive approach.

But after getting it working, look for repeated subtree calculations.

A strong follow-up thought process is:

Am I visiting the same subtree multiple times?
↓
Can I calculate its information once?
↓
Can I return that information to the parent?
↓
Postorder DFS

This is often the difference between an acceptable recursive solution and an optimal tree solution.

Conclusion

The straightforward approach is a natural way to solve this problem: for each node, calculate the sum and number of nodes in its subtree, find the average, and check whether it matches the node's value.

The first improvement is to notice that the same DFS used for calculating the sum can also maintain the node count using a shared counter instead of running a separate count() traversal.

However, the bigger optimization comes from noticing that subtrees are still being recalculated.

The optimal solution solves this by using one postorder DFS. Each node receives the sum and count from its children, calculates its own subtree information, checks the average, and passes the result to its parent.

This changes the worst-case time complexity from:

O(n²)

to:

O(n)

The most valuable takeaway is the pattern behind the solution:

When a tree problem asks for information about a subtree, try to calculate that information once and return it upward through postorder DFS.

Related Articles

LeetCode 144: Binary Tree Preorder Traversal – Java Recursive & Iterative Solution Explained

LeetCode 144: Binary Tree Preorder Traversal – Java Recursive & Iterative Solution Explained

Learn how to solve LeetCode 144 Binary Tree Preorder Traversal using Recursive and Iterative approaches in Java. Includes traversal intuition, stack-based solution, dry run, complexity analysis, interview explanation, and common mistakes.

Read Article
LeetCode 102: Binary Tree Level Order Traversal – Java BFS Solution Explained

LeetCode 102: Binary Tree Level Order Traversal – Java BFS Solution Explained

Learn how to solve LeetCode 102 Binary Tree Level Order Traversal using BFS and Queue in Java. Includes brute force intuition, optimized level order traversal, dry run, complexity analysis, interview explanation, and common mistakes.

Read Article
LeetCode 543: Diameter of Binary Tree – Java DFS Solution Explained

LeetCode 543: Diameter of Binary Tree – Java DFS Solution Explained

Learn how to solve LeetCode 543 Diameter of Binary Tree using optimized DFS recursion in Java. Includes brute force approach, optimized solution, intuition, dry run, complexity analysis, interview explanation, and common mistakes.

Read Article
LeetCode 145: Binary Tree Postorder Traversal – Java Recursive & Iterative Solution Explained

LeetCode 145: Binary Tree Postorder Traversal – Java Recursive & Iterative Solution Explained

Learn how to solve LeetCode 145 Binary Tree Postorder Traversal using Recursive and Iterative approaches in Java. Includes traversal intuition, stack-based solution, dry run, complexity analysis, interview explanation, and common mistakes.

Read Article
Ai Assistant Kas