Count Nodes Equal to Average of Subtree

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. A subtree of root is a tree consisting of root and all of its descendants. The average of n elements is the sum of the n elements divided by n and rounded down to the nearest integer.

For example:

Consider the binary tree: [4,8,5,0,1,null,6]

• 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.

Therefore, the expected output is 5.

Another example:

Consider the binary tree: [1]

• For the node with value 1: The average of its subtree is 1 / 1 = 1.

Therefore, the expected output is 1.

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: TreeNode) -> int:
        self.count = 0

        def dfs(node):
            if not node:
                return 0, 0  # (sum, count)

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

            total_sum = node.val + left_sum + right_sum
            total_count = 1 + left_count + right_count

            if node.val == total_sum // total_count:
                self.count += 1

            return total_sum, total_count

        dfs(root)
        return self.count

Explanation:

The averageOfSubtree function calculates the number of nodes in a binary tree where the node's value is equal to the average of its subtree. The approach uses a depth-first search (DFS) traversal.

1. TreeNode Definition: The code starts by defining the structure of a binary tree node, including val, left, and right pointers.
2. DFS Traversal: The dfs function recursively traverses the binary tree.
   • Base Case: If the node is None, it returns (0, 0) indicating a sum of 0 and a count of 0.
   • Recursive Calls: It recursively calls dfs on the left and right subtrees.
   • Sum and Count Calculation: It calculates the total sum of the subtree by adding the current node's value to the sums returned by the left and right subtree calls. Similarly, it computes the total node count by adding 1 (for the current node) to the counts returned by the left and right subtree calls.
   • Average Check: It checks if the current node's value is equal to the average of the subtree (total sum divided by total count, using integer division).
   • Increment Count: If the average equals the node's value, the self.count variable is incremented.
   • Return Value: The function returns the total sum and total count of the current subtree to its parent node.
3. Main Function: The averageOfSubtree function initializes a self.count variable to 0. It calls the dfs function starting from the root node and returns the final self.count.

Example:

For the input root = [4,8,5,0,1,null,6], the code will:

1. Start DFS from node 4.
2. Recursively traverse left and right subtrees.
3. Calculate subtree sums and counts at each node.
4. Check average condition:
   • Node 4: average is (4 + 8 + 5 + 0 + 1 + 6) // 6 = 4, count incremented.
   • Node 8: average is (8 + 0 + 1) // 3 = 3
   • Node 5: average is (5 + 6) // 2 = 5, count incremented.
   • Node 0: average is 0 // 1 = 0, count incremented.
   • Node 1: average is 1 // 1 = 1, count incremented.
   • Node 6: average is 6 // 1 = 6, count incremented.
5. Return final count of 5.

Edge Cases:

• Empty Tree: If the root is None, the function correctly returns 0 because dfs will not be called.
• Single Node Tree: If the tree consists of only the root node, the sum and count will be the node's value and 1 respectively. The average will equal the node's value, so count will increment to 1, which is correct.
• Nodes with Zero Values: The code correctly handles nodes with zero values since the sum and average calculations will work as expected.

Time and Space Complexity

Time Complexity: O(N), where N is the number of nodes in the binary tree. The DFS traversal visits each node once.

Space Complexity: O(H), where H is the height of the binary tree. This is due to the recursive call stack during the DFS traversal. In the worst-case scenario (skewed tree), H can be equal to N, resulting in O(N) space complexity. In the best-case scenario (balanced tree), H would be log(N), resulting in O(log(N)) space complexity.