Count Nodes Equals Leetcode 2265

// This is Brute Force solution 
// TC - O(N*N)
// SC - O(N)


class Solution {
    private int count ;
    private int sum ;
    public int averageOfSubtree(TreeNode root) {
        Queue<TreeNode> q = new LinkedList();

        q.offer(root);
        int ans = 0;
        while (!q.isEmpty()) {
            TreeNode current = q.poll();
            count = 0;
            sum = 0;
            helper(current) ;
            if ((sum/count) == current.val) {
                ans++;
            }

            if (current.left != null) {
                q.offer(current.left);
            }

            if (current.right != null) {
                q.offer(current.right);
            }

        }

        return ans;
    }

    private void helper(TreeNode root) {
        if (root == null) {
            return;
        }

        helper(root.left);
        helper(root.right);
        sum += root.val;
        count++;
    }
}


// This is a optimized version of this the question
// TC - O(N)
// SC - (N)  stack space


class Solution {
    private int ans ;
    public int averageOfSubtree(TreeNode root) {
        ans = 0;

        helper(root);

        return ans;
    }

    private int[] helper(TreeNode root) {
        if (root == null) {
            return new int[]{0, 0};
        }

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

        int count = 1 + left[0] + right[0];
        int sum = root.val + left[1] + right[1];
        if ((sum/count) == root.val) {
            ans++;
        }

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

}