Calculating the average of the entire tree separately and returning is incorrect - putting it out there because the problem can be misconstrued in that manner.
Instead, what’s asked is the running averages for every node, where you consider every node’s subtree and not the entire tree.
We return a list from our DFS function instead of a singular value since we need access to the node_count, running_sum, and current_answer in the same instances.
Code
Python
class Solution:
def averageOfSubtree(self, root: TreeNode) -> int:
def calc(node: TreeNode) -> List[int]:
if node == None: return [0,0,0]
left, right = calc(node.left), calc(node.right)
totalVal = left[0] + right[0] + node.val
count = left[1] + right[1] + 1
return [totalVal, count, left[2] + right[2] + (1 if totalVal // count == node.val else 0)]
return calc(root)[2]Big O Analysis
-
Runtime
The runtime complexity here is since we visit all nodes at least once.
-
Memory
The memory usage is since we would use the implicit call stack while making recursive tree calls.
— A