Count Nodes Equal to Average of Subtree

Sanjeev SharmaSanjeev Sharma
1 min read

Advertisement

Problem

Count nodes where node.val == average of the subtree rooted at that node (floor division).

Solutions

# Python
def averageOfSubtree(root):
    count = [0]
    def dfs(node):
        if not node:
            return 0, 0  # sum, count
        ls, lc = dfs(node.left)
        rs, rc = dfs(node.right)
        total_sum = ls + rs + node.val
        total_count = lc + rc + 1
        if total_sum // total_count == node.val:
            count[0] += 1
        return total_sum, total_count
    dfs(root)
    return count[0]
// C++
int cnt = 0;
pair<int,int> dfs(TreeNode* node) {
    if (!node) return {0,0};
    auto [ls,lc]=dfs(node->left), [rs,rc]=dfs(node->right);
    int s=ls+rs+node->val, c=lc+rc+1;
    if(s/c==node->val)cnt++;
    return {s,c};
}
int averageOfSubtree(TreeNode* root){cnt=0;dfs(root);return cnt;}
// Java
int cnt=0;
public int averageOfSubtree(TreeNode root){dfs(root);return cnt;}
int[] dfs(TreeNode node){
    if(node==null)return new int[]{0,0};
    int[]l=dfs(node.left),r=dfs(node.right);
    int s=l[0]+r[0]+node.val,c=l[1]+r[1]+1;
    if(s/c==node.val)cnt++;
    return new int[]{s,c};
}
// JavaScript
function averageOfSubtree(root) {
    let count = 0;
    function dfs(node) {
        if (!node) return [0, 0];
        const [ls,lc]=dfs(node.left), [rs,rc]=dfs(node.right);
        const s=ls+rs+node.val, c=lc+rc+1;
        if(Math.floor(s/c)===node.val)count++;
        return [s,c];
    }
    dfs(root);
    return count;
}

Complexity

  • Time: O(n), Space: O(h)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro