2265. Count Nodes Equal to Average of Subtree

solution

DFS

class Solution {
    int res = 0;
    public int averageOfSubtree(TreeNode root) {
        sum(root);
        return res;
    }

    // Returns [sum of sub tree, node counts of sub tree]
    public int[] sum(TreeNode node) {
        if (node == null) {
            return new int[]{0, 0};
        }
        int[] resLeft = sum(node.left);
        int[] resRight = sum(node.right);

        int leftSum = resLeft[0], leftCount = resLeft[1];
        int rightSum = resRight[0], rightCount = resRight[1];
        int sum = node.val + leftSum + rightSum;
        int count = 1 + leftCount + rightCount;
        int avg = sum / count;
        if (avg == node.val) {
            res++;
        }

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