Problem statement

https://binarysearch.com/problems/Subtree-with-Maximum-Average/

Solution

Equal to Leetcode 1120 Maximum Average Subtree.

Complexity

It is O(n) for time and O(h) for space.

Code

class Solution:
    def solve(self, root):
        self.ans = 0

        def dfs(node):
            if not node: return (0, 0)
            cnt1, sum1 = dfs(node.left)
            cnt2, sum2 = dfs(node.right)
            cnt3, sum3 = cnt1 + cnt2 + 1, sum1 + sum2 + node.val
            self.ans = max(self.ans, sum3/cnt3)
            return (cnt3, sum3)
        
        dfs(root)
        return self.ans