← All solutions
September 10Medium

Count nodes equal to average of subtree

  • Tree
  • DFS
  • Binary Tree

Problem

Given the root of a binary tree, return the number of nodes where the value of the node is equal to the average of the values in its subtree.

Note: The average of n elements is the sum of the n elements divided by n and rounded down to the nearest integer. A subtree of root is a tree consisting of root and all of its descendants.

Definition for a binary tree node. class TreeNode: def __init__(self, val=0, left=None, right=None): self.val = val self.left = left self.right = right

Solution

class Solution:
    def averageOfSubtree(self, root: TreeNode) -> int:
        if not root:
            return 0

        ans = 0
        def dfs(root):
            if not root:
                return 0, 0
            left = dfs(root.left)
            right = dfs(root.right)
            total = root.val + left[0] + right[0]
            count = 1 + left[1] + right[1]
            if floor(total / count) == root.val:
                nonlocal ans
                ans += 1
            return total, count

        dfs(root)
        return ans