Given a binary tree, find the subtree with maximum average. Return the root of the subtree.
Example 1
Input:
{1,-5,11,1,2,4,-2}
Output:11
Explanation:
The tree is look like this:
1
/ \
-5 11
/ \ / \
1 2 4 -2
The average of subtree of 11 is 4.3333, is the maximun.
Example 2
Input:
{1,-5,11}
Output:11
Explanation:
1
/ \
-5 11
The average of subtree of 1,-5,11 is 2.333,-5,11. So the subtree of 11 is the maximun.
"""
Definition of TreeNode:
class TreeNode:
def __init__(self, val):
self.val = val
self.left, self.right = None, None
"""
class Solution:
"""
@param root: the root of binary tree
@return: the root of the maximum average of subtree
"""
def findSubtree2(self, root):
def dfs(node):
"""
in:
node: treenode: root of current subtree
out:
maxnode: treenode: root of currently known
subtree with maximum average
maxmean: float: currently known max average
n: int: number of nodes in the current subtree
accu: int: sum of the nodes in the currently subtree
"""
# if no such node
if node is None:
return None, -float("inf"), 0, 0
# if leaf
if node.left is None and node.right is None:
return node, node.val, 1, node.val
# acquire answers
maxnodeL, maxmeanL, nL, accuL = dfs(node.left)
maxnodeR, maxmeanR, nR, accuR = dfs(node.right)
# combine answers
## get the average of the current tree
### find current $n$ and $accu$
n = nL + nR + 1
accu = accuL + accuR + node.val
### find average
mean = accu / n
## update the max record
maxmean = max([maxmeanL, maxmeanR, mean])
if maxmean == maxmeanL:
maxnode = maxnodeL
elif maxmean == maxmeanR:
maxnode = maxnodeR
else:
maxnode = node
# return combined answer
return maxnode, maxmean, n, accu
return dfs(root)[0]
Special care:
maxnode
maxmean
n
和accu
。