Lintcode 597. Subtree with Maximum Average

Given a binary tree, find the subtree with maximum average. Return the root of the subtree.

Example

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.

Notice

LintCode will print the subtree which root is your return node.
It's guaranteed that there is only one subtree with maximum average.

Code (Java):

/**
 * Definition of TreeNode:
 * public class TreeNode {
 *     public int val;
 *     public TreeNode left, right;
 *     public TreeNode(int val) {
 *         this.val = val;
 *         this.left = this.right = null;
 *     }
 * }
 */

public class Solution {
    /**
     * @param root: the root of binary tree
     * @return: the root of the maximum average of subtree
     */
    public TreeNode findSubtree2(TreeNode root) {
        // write your code here
        if (root == null) {
            return null;
        }
        
        ResultType ans = findSubtree2Helper(root);
        
        return ans.node;
    }
    
    private ResultType findSubtree2Helper(TreeNode root) {
        if (root == null) {
            return new ResultType(0, 0, Integer.MIN_VALUE, null);
        }
        
        ResultType left = findSubtree2Helper(root.left);
        ResultType right = findSubtree2Helper(root.right);
        
        int sum = root.val + left.sum + right.sum;
        int numNodes = left.numNodes + right.numNodes + 1;
        double ave = (double) sum / numNodes;
        TreeNode node = null;
        if (ave > left.maxAve && ave > right.maxAve) {
            node = root;
        } else if (left.maxAve > ave && left.maxAve > right.maxAve) {
            ave = left.maxAve;
            node = left.node;
        }  else {
            ave = right.maxAve;
            node = right.node;
        }
        
        return new ResultType(sum, numNodes, ave, node);
    }
}

class ResultType {
    int sum;
    int numNodes;
    double maxAve;
    TreeNode node;
    
    public ResultType(int sum, int numNodes, double maxAve, TreeNode node) {
        this.sum = sum;
        this.numNodes = numNodes;
        this.maxAve = maxAve;
        this.node = node;
    }
}