Leet code problem 2265: Count Nodes Equal to Average of Subtree

This code defines a class called Solution with a method named averageOfSubtree. The method takes a binary tree's root node as input and calculates the count of nodes where the average value of the subtree rooted at that node is equal to the value of the node itself.

Here's a breakdown of the code:

ans is initialized as an empty list. This list will store the values of nodes where the average of their subtree is equal to their own value.

The dfs function is defined as a nested function within the averageOfSubtree method. It is a depth-first search function used to traverse the binary tree recursively. It takes a root node as input and returns a tuple of two values: the sum of all node values in the subtree rooted at root, and the count of nodes in that subtree.

Inside the dfs function, it checks if the root is None, indicating an empty subtree. If so, it returns (0, 0) to represent the sum and count of an empty subtree.

If the root is not None, it recursively calls dfs on the left and right subtrees (root.left and root.right). It then calculates leftSum and leftCount for the left subtree and rightSum and rightCount for the right subtree.

The total sum is calculated as the sum of the current node's value, the leftSum, and the rightSum.

The total count is calculated as 1 (for the current node) plus the leftCount and rightCount.

It checks whether the average of the subtree (sum // count) is equal to the value of the current node. If it is, it appends the current node's value to the ans list.

The dfs function returns a tuple containing the calculated sum and count for the current subtree.

The dfs function is called initially with the root of the binary tree to start the recursive traversal.

Finally, the averageOfSubtree method returns the length of the ans list, which represents the count of nodes whose subtree averages are equal to the node values.

In essence, this code calculates the number of nodes in a binary tree where the average of the values in their subtrees matches their own value and stores these values in the ans list. The length of this list is returned as the final result.

Code in python

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

        def dfs(root: Optional[TreeNode]) -> Tuple[int, int]:
            nonlocal ans
            if not root:
                return (0, 0)
            leftSum, leftCount = dfs(root.left)
            rightSum, rightCount = dfs(root.right)
            sum = root.val + leftSum + rightSum
            count = 1 + leftCount + rightCount
            if sum // count == root.val:
                ans.append(root.val)
            return (sum, count)

        dfs(root)
        return len(ans)

Code in Ruby

# Definition for a binary tree node.
# class TreeNode
#     attr_accessor :val, :left, :right
#     def initialize(val = 0, left = nil, right = nil)
#         @val = val
#         @left = left
#         @right = right
#     end
# end
# @param {TreeNode} root
# @return {Integer}
def average_of_subtree(root)
    @ans = []

    def dfs(root)
      return [0, 0] if root.nil?

      left_sum, left_count = dfs(root.left)
      right_sum, right_count = dfs(root.right)
      summ = root.val + left_sum + right_sum
      count = 1 + left_count + right_count
      @ans << root.val if summ / count == root.val
      return [summ, count]
    end

    dfs(root)
    @ans.length

end

Code in Javascript

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 * @return {number}
 */
function averageOfSubtree(root) {
    let ans = [];
    function dfs(root) {
        if (!root) {
            return [0, 0];
        }
        let [leftSum, leftCount] = dfs(root.left);
        let [rightSum, rightCount] = dfs(root.right);
        let sum = root.val + leftSum + rightSum;
        let count = 1 + leftCount + rightCount;
        if (Math.floor(sum / count) === root.val) {
            ans.push(root.val);
        }
        return [sum, count];
    }
    dfs(root);
    return ans.length;
}