2265. Count Nodes Equal to Average of Subtree

Tags:

• Medium
• dfs
• tree

Link

Get the (nodeCount, sum) from the bottom.

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

        if not root:
            return 0

        self.res = 0

        def dfs(node):
            if not node:
                return (0, 0)

            left = dfs(node.left)
            right = dfs(node.right)

            s = left[1] + right[1] + node.val

            avg = s // (1 + left[0] + right[0])

            if avg == node.val:

                self.res += 1

            return (left[0] + right[0] + 1, s)

        dfs(root)

        return self.res