# 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
class Solution:
def averageOfSubtree(self, root: Optional[TreeNode]) -> int:
# return 3 values
# totalSum of the subtree, totoalNodes, numNodes that value==avg
def dfs(node):
if not node:
return 0, 0, 0
totalSumLeft, totalNodesLeft, resLeft = dfs(node.left)
totalSumRight, totalNodesRight, resRight = dfs(node.right)
totalSum = totalSumLeft + totalSumRight + node.val
totalNodes = totalNodesLeft + totalNodesRight + 1
res = resLeft + resRight
if node.val == totalSum // totalNodes:
res += 1
return totalSum, totalNodes, res
_, _, res = dfs(root)
return res
