2265. Count Nodes Equal to Average of Subtree - Leetcode Solution

Leetcode Problem Link Code Solution Link:

💡 Step-by-Step Thought Process

Understand the problem: Count nodes in a binary tree where the node’s value equals the average of its subtree’s node values (integer division).
Initialize a list (num_nodes) with a single element 0 to count matching nodes.
Define a DFS function that returns a tuple of the number of nodes and sum of values in a subtree.
If the root is None, return (0, 0).
Recursively compute the number of nodes (N_left) and sum (summ_left) for the left subtree.
Compute the same for the right subtree (N_right, summ_right).
Calculate the total nodes (N) as 1 plus N_left plus N_right and total sum (summ) as root value plus summ_left plus summ_right.
Compute the average as summ divided by N (integer division).
If the root’s value equals the average, increment num_nodes[0].
Return (N, summ).
Call DFS on the root and return num_nodes[0].

Code Solution

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

        def dfs(root):
            if not root:
                return (0, 0)
            
            N_left, summ_left = dfs(root.left)
            N_right, summ_right = dfs(root.right)

            N = 1 + N_left + N_right
            summ = root.val + summ_left + summ_right
            avg = summ // N

            if root.val == avg:
                num_nodes[0] += 1

            return (N, summ)
        
        dfs(root)
        return num_nodes[0]
        # Time: O(N)
        # Space: O(N)

#include <algorithm>
using namespace std;

class Solution {
private:
    int numNodes = 0;

public:
    int averageOfSubtree(TreeNode* root) {
        dfs(root);
        return numNodes;
    }

    // Returns a pair: {number of nodes, sum of values}
    pair<int, int> dfs(TreeNode* root) {
        if (!root) return {0, 0};

        auto left = dfs(root->left);
        auto right = dfs(root->right);

        int nodeCount = 1 + left.first + right.first;
        int sum = root->val + left.second + right.second;
        int avg = sum / nodeCount;

        if (root->val == avg) numNodes++;

        return {nodeCount, sum};
    }
};

public class Solution {
    private int numNodes = 0;

    public int averageOfSubtree(TreeNode root) {
        dfs(root);
        return numNodes;
    }

    private int[] dfs(TreeNode root) {
        if (root == null) return new int[] {0, 0}; // {node count, sum of values}

        int[] left = dfs(root.left);
        int[] right = dfs(root.right);

        int nodeCount = 1 + left[0] + right[0];
        int sum = root.val + left[1] + right[1];
        int avg = sum / nodeCount;

        if (root.val == avg) numNodes++;

        return new int[] {nodeCount, sum};
    }
}

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = val;
 *     this.left = left;
 *     this.right = right;
 * }
 */

/**
 * @param {TreeNode} root
 * @return {number}
 */
var averageOfSubtree = function(root) {
    let numNodes = 0;

    function dfs(node) {
        if (!node) return [0, 0]; // [node count, sum of values]

        let [leftCount, leftSum] = dfs(node.left);
        let [rightCount, rightSum] = dfs(node.right);

        let nodeCount = 1 + leftCount + rightCount;
        let sum = node.val + leftSum + rightSum;
        let avg = Math.floor(sum / nodeCount);

        if (node.val === avg) numNodes++;

        return [nodeCount, sum];
    }

    dfs(root);
    return numNodes;
};

Detailed Explanation

Problem Overview

The "Count Nodes Equal to Average of Subtree" problem focuses on binary trees. You are given a binary tree, and your task is to count how many nodes have a value equal to the average of all nodes in their respective subtree (including themselves). The average must be calculated using integer division. This question tests your understanding of depth-first search (DFS), tree traversal, and recursive aggregation of node data.

Key Insights and Approach

The most important observation is that at each node, we need to compute two things for its entire subtree: the total number of nodes and the total sum of those nodes’ values. With this information, we can easily compute the average for the subtree and check whether it equals the value of the root node of that subtree.

A recursive depth-first search (DFS) approach is ideal here. As we traverse the tree from the bottom up, each recursive call can return both the number of nodes and the sum of values in the subtree rooted at the current node. Once we have these values, we check whether the current node’s value equals the computed average using integer division.

Step-by-Step Strategy

To solve this efficiently, we define a recursive helper function that visits each node and returns a tuple: (count_of_nodes_in_subtree, sum_of_subtree_values).

For each node:

Recursively compute subtree size and value sum for the left child.
Do the same for the right child.
Combine the counts and sums from both children with the current node’s value.
Calculate the average by dividing the total sum by the total count (using integer division).
Check if the current node’s value matches this average. If it does, increment a global counter.

Time and Space Complexity

Time Complexity: O(n), where n is the number of nodes in the binary tree. Each node is visited exactly once, and all necessary computations are performed in constant time per node.
Space Complexity: O(h), where h is the height of the tree, due to the recursion stack. In the worst case (a skewed tree), this can be O(n).

Why This Approach Works

This method is both elegant and efficient because it leverages the natural structure of the tree and avoids unnecessary recomputation. By gathering data in a post-order fashion (left, right, root), we ensure that each subtree is fully processed before its result is used in a parent’s calculation. This enables a one-pass solution without auxiliary data structures or repeated traversals.

Conclusion

Counting nodes equal to the average of their subtree reinforces recursive thinking and showcases how to aggregate and compare subtree information efficiently. This problem also offers a practical example of combining multiple return values in recursive tree functions, which is a useful technique in more complex tree-based algorithms.

Get Personalized Support at AlgoMap Bootcamp 💡

AlgoMap.io

Want Help Cracking FAANG?