Approach: Post-order returning (sum, count) of each subtree

def averageOfSubtree(self, root):
    self.cnt = 0
    def dfs(node):
        if not node: return (0, 0)
        ls, lc = dfs(node.left); rs, rc = dfs(node.right)
        s = ls + rs + node.val; c = lc + rc + 1
        if s // c == node.val: self.cnt += 1
        return (s, c)
    dfs(root)
    return self.cnt

Complexity: O(n) time, O(h) space.

Full Solution — all languages

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):
        self.cnt = 0
        def dfs(node):
            if not node: return (0, 0)
            ls, lc = dfs(node.left); rs, rc = dfs(node.right)
            s = ls + rs + node.val; c = lc + rc + 1
            if s // c == node.val: self.cnt += 1
            return (s, c)
        dfs(root)
        return self.cnt

import java.util.*;
class TreeNode {
    int val;
    TreeNode left, right;
    TreeNode(int v) {
        val = v;
    }
}
class Solution {
    public int averageOfSubtree(TreeNode root) {
        cnt = 0;
        dfs(root);
        return cnt;
    }
    int cnt;
    private 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};
    }
}

#include <bits/stdc++.h> using namespace std;
struct TreeNode {
    int val;
    TreeNode *left, *right;
    TreeNode(int v): val(v), left(nullptr), right(nullptr) {
    }
};
class Solution {
    public: int averageOfSubtree(TreeNode* root) {
        int cnt = 0;
        function<pair<int,int>(TreeNode*)> dfs = [&](TreeNode* node) -> pair<int,int> {
            if (!node) return {
                0, 0
            };
            auto l = dfs(node->left);
            auto r = dfs(node->right);
            int s = l.first + r.first + node->val, c = l.second + r.second + 1;
            if (s / c == node->val) cnt++;
            return {
                s, c
            };
        };
        dfs(root);
        return cnt;
    }
};

/**
 * Tree node shape: { val, left, right }
 * @param {TreeNode} root
 * @return {number}
 */ var averageOfSubtree = function(root) {
    let cnt = 0;
    const dfs = (node) => {
        if (!node) return [0, 0];
        const l = dfs(node.left), r = dfs(node.right);
        const s = l[0] + r[0] + node.val, c = l[1] + r[1] + 1;
        if (Math.floor(s / c) === node.val) cnt++;
        return [s, c];
    };
    dfs(root);
    return cnt;
};

Count Nodes Equal to Average of Subtree

Example 1

Example 2

Constraints