LeetCode 2265 - Count Nodes Equal to Average of Subtree

LeetCode Problem 2265

Difficulty: 🟡 Medium
Topics: Tree, Depth-First Search, Binary Tree

Solution

Problem Understanding

This problem asks us to count how many nodes in a binary tree satisfy a specific condition: the node’s value must equal the average value of all nodes in its subtree.

A subtree consists of the current node and every descendant below it. For each node, we need to compute:

• The total sum of all values in its subtree
• The total number of nodes in that subtree
• The integer average, computed using floor division

If the node’s value equals this computed average, we count it.

For example, suppose a subtree contains values [5, 6]. The sum is 11, the number of nodes is 2, and the average is:

$$11 // 2 = 5$$

Since integer division rounds down, node 5 would qualify.

The input is the root of a binary tree. Each node contains an integer value and pointers to its left and right children. The output is a single integer representing the number of nodes whose values equal the average of their respective subtrees.

The constraints tell us that the tree contains at most 1000 nodes and each node value is between 0 and 1000.

Although 1000 nodes is not extremely large, a naive approach that repeatedly traverses subtrees can still become inefficient because many nodes would be visited multiple times. This suggests that we should look for a way to compute subtree information only once.

There are several important edge cases to keep in mind:

• A single-node tree always qualifies because its subtree consists only of itself, so its average equals its value.
• Leaf nodes always qualify for the same reason, since the subtree size is 1.
• Trees with many overlapping subtrees can make repeated calculations expensive if we recompute subtree sums from scratch for every node.
• Integer division matters. Since the problem explicitly states the average is rounded down, we must use floor division (// in Python, integer division in Go).

Approaches

Brute Force Approach

A straightforward way to solve this problem is to process every node independently.

For each node, we can recursively traverse its entire subtree to compute two things:

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

After computing these values, we calculate:

$$\text{average} = \text{sum} // \text{count}$$

If the average equals the node’s value, we increment our answer.

This method is correct because every node explicitly computes the exact statistics of its subtree.

However, the inefficiency comes from repeated work. Suppose a node appears in many ancestor subtrees. Its value gets recalculated repeatedly. In a skewed tree of n nodes, the first node traverses n nodes, the next traverses n - 1, and so on, resulting in roughly:

$$O(n^2)$$

time complexity.

Optimal Approach, Postorder DFS

The key observation is that subtree information can be computed bottom-up.

To determine whether a node matches its subtree average, we only need:

• The sum of the left subtree
• The count of nodes in the left subtree
• The sum of the right subtree
• The count of nodes in the right subtree

This naturally suggests a postorder traversal, where we process children before their parent.

For every node, we recursively gather:

• Total subtree sum
• Total subtree size

Then we compute the average and immediately check whether the node qualifies.

Since every node is processed exactly once, we eliminate redundant work.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n²)	O(h)	Recomputes subtree sums repeatedly
Optimal	O(n)	O(h)	Postorder DFS computes subtree data once

Here, h represents the tree height.

Algorithm Walkthrough

Optimal Algorithm, Postorder DFS

1. Start a depth-first traversal from the root using recursion.
2. For each node, recursively process the left subtree first. The recursive call returns two values:

• The sum of values in the left subtree
• The number of nodes in the left subtree

3. Recursively process the right subtree in the same way.
4. Combine the subtree information:

• Current subtree sum = left sum + right sum + current node value
• Current subtree count = left count + right count + 1

5. Compute the subtree average using integer division:

$$\text{average} = \text{subtree sum} // \text{subtree count}$$ 6. Compare the average with the current node value. If they are equal, increment the answer. 7. Return the subtree sum and subtree count to the parent node. 8. Continue until the recursion finishes at the root.

Why it works

The algorithm works because of the postorder traversal invariant:

When processing a node, the algorithm already knows the exact sum and size of both child subtrees.

Since a subtree consists of the current node plus its descendants, combining the child results gives the exact subtree statistics for the current node. Therefore, every average calculation is correct, and every node is evaluated exactly once.

Python Solution

# 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

from typing import Optional, Tuple

class Solution:
    def averageOfSubtree(self, root: Optional[TreeNode]) -> int:
        matching_nodes = 0

        def dfs(node: Optional[TreeNode]) -> Tuple[int, int]:
            nonlocal matching_nodes

            if not node:
                return 0, 0

            left_sum, left_count = dfs(node.left)
            right_sum, right_count = dfs(node.right)

            subtree_sum = left_sum + right_sum + node.val
            subtree_count = left_count + right_count + 1

            subtree_average = subtree_sum // subtree_count

            if subtree_average == node.val:
                matching_nodes += 1

            return subtree_sum, subtree_count

        dfs(root)
        return matching_nodes

The implementation follows a recursive postorder traversal.

The nested dfs() function returns a tuple containing the subtree sum and subtree node count. This makes it easy for parent nodes to combine results from their children.

The base case handles None nodes by returning (0, 0), meaning an empty subtree contributes no sum and no nodes.

After recursively processing the left and right children, the algorithm computes the current subtree statistics. It then calculates the average using integer division and compares it with the current node value.

The matching_nodes variable is declared outside the DFS function and updated using nonlocal, allowing every recursive call to contribute to the final answer.

Finally, the traversal starts at the root, and the accumulated count is returned.

Go Solution

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func averageOfSubtree(root *TreeNode) int {
    matchingNodes := 0

    var dfs func(node *TreeNode) (int, int)

    dfs = func(node *TreeNode) (int, int) {
        if node == nil {
            return 0, 0
        }

        leftSum, leftCount := dfs(node.Left)
        rightSum, rightCount := dfs(node.Right)

        subtreeSum := leftSum + rightSum + node.Val
        subtreeCount := leftCount + rightCount + 1

        subtreeAverage := subtreeSum / subtreeCount

        if subtreeAverage == node.Val {
            matchingNodes++
        }

        return subtreeSum, subtreeCount
    }

    dfs(root)
    return matchingNodes
}

The Go implementation closely mirrors the Python version.

Instead of a nested function with nonlocal, Go captures matchingNodes through a closure. Since Go uses integer division by default for integers, no extra rounding logic is needed.

Nil nodes are handled explicitly with:

if node == nil

which returns (0, 0) for empty subtrees.

Since the maximum possible subtree sum is only 1000 * 1000 = 1,000,000, integer overflow is not a concern with Go’s int type.

Worked Examples

Example 1

Input:

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

Tree:

        4
       / \
      8   5
     / \   \
    0   1   6

Because we use postorder traversal, children are processed before parents.

Node	Left Sum	Left Count	Right Sum	Right Count	Subtree Sum	Subtree Count	Average	Match?
0	0	0	0	0	0	1	0	Yes
1	0	0	0	0	1	1	1	Yes
8	0	1	1	1	9	3	3	No
6	0	0	0	0	6	1	6	Yes
5	0	0	6	1	11	2	5	Yes
4	9	3	11	2	24	6	4	Yes

Final count:

5

Matching nodes are:

0, 1, 5, 6, 4

Example 2

Input:

root = [1]

Tree:

1

Node	Left Sum	Right Sum	Subtree Sum	Count	Average	Match?
1	0	0	1	1	1	Yes

Final answer:

1

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Each node is visited exactly once
Space	O(h)	Recursion stack depends on tree height

The time complexity is O(n) because every node participates in exactly one DFS computation. No subtree is recomputed.

The space complexity is O(h), where h is the height of the tree due to recursive call stack usage. In a balanced tree this is O(log n), while in a completely skewed tree it becomes O(n).

Test Cases

# Helper class for testing
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

solution = Solution()

# Example 1
root1 = TreeNode(
    4,
    TreeNode(8, TreeNode(0), TreeNode(1)),
    TreeNode(5, None, TreeNode(6))
)
assert solution.averageOfSubtree(root1) == 5  # Provided example

# Example 2
root2 = TreeNode(1)
assert solution.averageOfSubtree(root2) == 1  # Single node tree

# All nodes match
root3 = TreeNode(
    0,
    TreeNode(0),
    TreeNode(0)
)
assert solution.averageOfSubtree(root3) == 3  # Every node equals average

# No internal match
root4 = TreeNode(
    10,
    TreeNode(1),
    TreeNode(1)
)
assert solution.averageOfSubtree(root4) == 2  # Only leaf nodes match

# Left skewed tree
root5 = TreeNode(
    3,
    TreeNode(
        2,
        TreeNode(1)
    )
)
assert solution.averageOfSubtree(root5) == 2  # Leaves plus valid averages

# Right skewed tree
root6 = TreeNode(
    5,
    None,
    TreeNode(
        5,
        None,
        TreeNode(5)
    )
)
assert solution.averageOfSubtree(root6) == 3  # Repeated values

# Integer division behavior
root7 = TreeNode(
    2,
    TreeNode(1),
    TreeNode(3)
)
assert solution.averageOfSubtree(root7) == 3  # Root average = 2

# Large equal-value tree
root8 = TreeNode(
    7,
    TreeNode(7, TreeNode(7), TreeNode(7)),
    TreeNode(7)
)
assert solution.averageOfSubtree(root8) == 6  # Every node matches

Test	Why
[4,8,5,0,1,null,6]	Validates the provided example
[1]	Tests single-node tree
[0,0,0]	Ensures all nodes can match
[10,1,1]	Tests internal node mismatch
Left skewed tree	Validates recursion depth and subtree aggregation
Right skewed tree	Ensures asymmetric trees work correctly
[2,1,3]	Verifies integer division behavior
Equal-value tree	Confirms all nodes can satisfy the condition

Edge Cases

Single Node Tree

A tree with only one node is the smallest valid input. Since the subtree contains only the node itself, the sum equals the node value and the count is 1. The average must therefore equal the node value. A buggy implementation might accidentally treat missing children incorrectly or divide by zero, but our implementation safely returns (0, 0) for null children and computes the average correctly.

Leaf Nodes

Every leaf node always qualifies because its subtree contains only itself. This can be easy to overlook when implementing recursion. Since our DFS returns subtree information bottom-up, a leaf naturally computes:

sum = node.val
count = 1
average = node.val

making the condition automatically correct.

Integer Division Rounding

The problem explicitly requires averages to be rounded down. For example, if the subtree sum is 11 and the count is 2, the average should be 5, not 5.5 or 6. A common mistake is to use floating-point division and compare decimals. Our implementation avoids this issue entirely by using integer division:

subtree_sum // subtree_count

and:

subtreeSum / subtreeCount

which naturally perform floor division.

Highly Skewed Trees

A tree may degenerate into a linked list, where every node has only one child. In such cases, recursion depth becomes O(n) instead of O(log n). Our algorithm still remains correct because every node is processed once, though stack usage increases proportionally to the tree height.