【Lintcode】597. Subtree with Maximum Average

题目地址：

https://www.lintcode.com/problem/subtree-with-maximum-average/description

给定一棵二叉树，求其子树使得该子树的平均值最大。返回子树的根。

由于要找平均值，所以我们在递归的时候需要同时把左右子树的节点数和数字和都返回给上一层。同时我们可以用两个全局变量来记录已经找到的最大平均值和对应的树根。最后返回那个树根即可。代码如下：

public class Solution {
    // 注意最小的Double应该是负的Double.MAX_VALUE；
    // Double.MIN_VALUE实际上只是个非常小的正数
    double maxAve = -Double.MAX_VALUE;
    TreeNode ans = null;
    /**
     * @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 root;
        }
        
        dfs(root);
        return ans;
    }
    
    private int[] dfs(TreeNode root) {
    	int[] res = new int[2];
        if (root == null) {
            return res;
        }
        // 分别计算左右子树的节点数和数字和
        int[] left = dfs(root.left), right = dfs(root.right);
        // 当前树的节点数就是左右子树的节点数之和 + 1；数字和就是左右子树数字之和 + 树根
        res[0] = left[0] + right[0] + 1;
        res[1] = left[1] + right[1] + root.val;
        // 算一下当前子树的平均值，如果更大，则更新全局变量
        double ave = (double) res[1] / res[0];
        if (ave > maxAve) {
            maxAve = ave;
            ans = root;
        }
        
        return res;
    }
}

class TreeNode {
    int val;
    TreeNode left, right;
    TreeNode(int x) {
        val = x;
    }
}

时空复杂度 $O(n)$。