Subtree with Maximum Average

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

LintCode will print the subtree which root is your return node.

It's guaranteed that there is only one subtree with maximum average.

Code:

/**
 * 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 {

    private TreeNode subtree = null;
    private ResultType subtreeResult = null;

    public TreeNode findSubtree2(TreeNode root) {
        if (root == null) {
            return null;
        }
        helper(root);
        return subtree;
    }

    private ResultType helper(TreeNode root) {
        if (root == null) {
            return new ResultType(0, 0);
        }
        ResultType left = helper(root.left);
        ResultType right = helper(root.right);
        ResultType current = new ResultType(
            left.sum + right.sum + root.val,
            left.size + right.size + 1
        );
        if (subtree == null || subtreeResult.sum * current.size < subtreeResult.size * current.sum) {
            subtree = root;
            subtreeResult = current;
        }
        return current;
    }

    class ResultType {
        int sum;
        int size;
        public ResultType(int sum, int size) {
            this.sum = sum;
            this.size = size;
        }
    }
}