[LintCode] 597. Subtree with Maximum Average

Posted 2020-11-27 xuanlu

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.

public class Solution {
    private class ResType {
        int sum;
        int size;
        public ResType(int sum, int size) {
            this.sum = sum;
            this.size = size;
        }
    }
    /**
     * @param root: the root of binary tree
     * @return: the root of the maximum average of subtree
     */
    TreeNode maxNode = null;
    ResType globalNode = null;
    public TreeNode findSubtree2(TreeNode root) {
        // write your code here
        helper(root);
        return maxNode;
    }
    
    private ResType helper(TreeNode root) {
        if (root == null) {
            return new ResType(0, 0);
        }
        ResType left = helper(root.left);
        ResType right = helper(root.right);
        int curSum = left.sum + right.sum + root.val;
        int curSize = left.size + right.size + 1;
        ResType cur = new ResType(curSum, curSize);
        if (maxNode == null || curSum * globalNode.size > curSize * globalNode.sum) {
            globalNode = cur;
            maxNode = root;
        }
        return cur;
    }
}