637. 二叉树的层平均值

题目：

给定一个非空二叉树, 返回一个由每层节点平均值组成的数组.

示例 1:

输入:
    3
   / \
  9  20
    /  \
   15   7
输出: [3, 14.5, 11]
解释:
第0层的平均值是 3,  第1层是 14.5, 第2层是 11. 因此返回 [3, 14.5, 11].
注意：

节点值的范围在32位有符号整数范围内。

思路：

我们可以使用深度优先搜索遍历整颗二叉树。我们使用两个数组 sum 存放树中每一层的节点数值之和，以及 count 存放树中每一层的节点数量之和。在遍历时，我们需要额外记录当前节点所在的高度，并根据高度 h 更新数组元素 sum[h] 和 count[h]。在遍历结束之后，res = sum / count 即为答案。

代码：

import java.util.ArrayList;
import java.util.List;

/**
 * 637. 二叉树的层平均值
 * @author linwillen
 * @create 2020-06-20-13:48
 */
public class AverageOfLevels {

    public List<Double> averageOfLevels(TreeNode root) {
        List<Double> sum = new ArrayList<>();
        List<Integer> count = new ArrayList<>();
        List<Double> res = new ArrayList<>();
        average(root,0,sum,count);
        for (int i = 0;i<sum.size();i++){
            res.add(sum.get(i)/count.get(i));
        }
        return res;
    }

    /**
     *
     * @param t 表示树的当前节点
     * @param i 表示当前树的高度
     * @param sum   表示当前高度的总和
     * @param count 表示当前高度节点的数量
     */
    public void average(TreeNode t, int i, List<Double> sum, List<Integer> count) {
        if (t == null) {
            return;
        }
        if (i < sum.size()) {
            sum.set(i,sum.get(i)+t.val);
            count.set(i,count.get(i)+1);
        }else{
            sum.add((double)t.val);
            count.add(1);
        }
        average(t.left, i + 1, sum, count);
        average(t.right, i + 1, sum, count);
    }
}