# Count Nodes Equal to Average of Subtree
**Difficulty:** MEDIUM
[External](https://leetcode.com/problems/count-nodes-equal-to-average-of-subtree)
Canonical: https://scaleengineer.com/dsa/problems/count-nodes-equal-to-average-of-subtree
**Algorithms:** [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search)
**Data structures:** Tree, Binary Tree
**Companies:** [Snowflake](https://scaleengineer.com/companies/snowflake), [INDmoney](https://scaleengineer.com/companies/indmoney)
---
## Problem
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:**

![](https://assets.glich.co/dsa/count-nodes-equal-to-average-of-subtree/image0.png) 

**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:**

![](https://assets.glich.co/dsa/count-nodes-equal-to-average-of-subtree/image1.png) 

**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
## Brute Force with Nested Traversal
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.
**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.
**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.
### Explanation
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.

```java
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    int result = 0;

    public int averageOfSubtree(TreeNode root) {
        // Main traversal to visit every node
        traverse(root);
        return result;
    }

    // This function iterates through each node of the tree
    private void traverse(TreeNode node) {
        if (node == null) {
            return;
        }

        // For the current node, get its subtree's sum and count
        int[] subtreeInfo = getSubtreeInfo(node);
        int sum = subtreeInfo[0];
        int count = subtreeInfo[1];

        // Check if the node's value equals the average of its subtree
        if (count > 0 && (sum / count) == node.val) {
            result++;
        }

        // Continue to the next nodes
        traverse(node.left);
        traverse(node.right);
    }

    // This helper function calculates the sum and count for a given subtree
    private int[] getSubtreeInfo(TreeNode subRoot) {
        if (subRoot == null) {
            return new int[]{0, 0}; // {sum, count}
        }

        int[] leftInfo = getSubtreeInfo(subRoot.left);
        int[] rightInfo = getSubtreeInfo(subRoot.right);

        int totalSum = subRoot.val + leftInfo[0] + rightInfo[0];
        int nodeCount = 1 + leftInfo[1] + rightInfo[1];

        return new int[]{totalSum, nodeCount};
    }
}
```
### 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`.

## Optimized Post-order Traversal
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.
**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).
**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).
### Explanation
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.

```java
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    int resultCount = 0;

    public int averageOfSubtree(TreeNode root) {
        postOrder(root);
        return resultCount;
    }

    // Returns a pair [sum, count] for the subtree rooted at 'node'
    private int[] postOrder(TreeNode node) {
        if (node == null) {
            return new int[]{0, 0}; // {sum, count}
        }

        // 1. Traverse left and right subtrees first (post-order)
        int[] leftSubtree = postOrder(node.left);
        int[] rightSubtree = postOrder(node.right);

        // 2. Process the current node using results from children
        int currentSum = node.val + leftSubtree[0] + rightSubtree[0];
        int currentCount = 1 + leftSubtree[1] + rightSubtree[1];

        // Calculate average and check the condition
        if (currentSum / currentCount == node.val) {
            resultCount++;
        }

        // 3. Return sum and count for the parent's calculation
        return new int[]{currentSum, currentCount};
    }
}
```
### 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`.

# Solutions
### Java

```java
/** * 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 }; } }
```

### CPP

```cpp
/** * 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 }; } };
```

### Python

```python
# 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
```