LeetCode Solutions Blog

Count Nodes Equal to Average of Subtree

Number: 2347

Difficulty: Medium

Paid? No

Problem Description

Given the root of a binary tree, determine the number of nodes such that the value of the node is equal to the average of the values in its subtree. The average is computed by summing all the values in the subtree (including the node itself), dividing by the number of nodes in the subtree, and rounding down to the nearest integer.

Key Insights

Use depth-first search (DFS) to traverse the tree in a post-order fashion.
For each node, compute the cumulative sum and count of nodes in its subtree.
Calculate the average as the floor division of the sum by the count.
Check if the node's value matches the computed average and increment a global counter accordingly.

Space and Time Complexity

Time Complexity: O(n), where n is the number of nodes in the tree since every node is visited once.
Space Complexity: O(h), where h is the height of the tree, due to recursion stack space (worst-case O(n) for a skewed tree).

Solution

Perform a post-order DFS traversal. For every node:

Recursively obtain the sum and count of nodes from the left and right subtrees.
Calculate the current subtree’s sum and total node count.
Compute the average (using integer division which rounds down).
Compare the node’s value with the calculated average and update the counter if they match.
Return the pair (subtree sum, subtree count) to be used for computing the averages in upper recursive calls.

Code Solutions

# Define the structure for the tree node
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

# Solution class to compute the count of valid nodes
class Solution:
    def averageOfSubtree(self, root: TreeNode) -> int:
        self.result = 0  # Global counter for valid nodes
        
        def dfs(node):
            # Base case: if node is None, return sum=0 and count=0
            if not node:
                return (0, 0)
            # Recursively get the sum and count from the left subtree
            left_sum, left_count = dfs(node.left)
            # Recursively get the sum and count from the right subtree
            right_sum, right_count = dfs(node.right)
            
            # Compute current subtree's total sum and count (including current node)
            subtree_sum = node.val + left_sum + right_sum
            subtree_count = 1 + left_count + right_count
            
            # Compute the average (floor division)
            average = subtree_sum // subtree_count
            
            # If the node's value equals the computed average, increment the result
            if average == node.val:
                self.result += 1
            
            # Return the total sum and count for this subtree
            return (subtree_sum, subtree_count)
        
        dfs(root)
        return self.result