LeetCode 637.二叉树的层平均值

问题描述

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

示例 1：

输入：
    3
   / \
  9  20
    /  \
   15   7
输出：[3, 14.5, 11]

解释：
第 0 层的平均值是 3 , 第1层是 14.5 , 第2层是 11 。因此返回 [3, 14.5, 11] 。

提示：

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

解题示例1

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public List<Double> averageOfLevels(TreeNode root) {
        // 记录每层节点和
        List<Double> sums = new ArrayList<Double>();
        //记录每层节点个数
        List<Integer> counts = new ArrayList<Integer>();
        // 返回结果
        List<Double> res = new ArrayList<Double>();

        dfs(root, 0, counts, sums);
        for(int i = 0; i < sums.size(); ++i) {
            res.add(sums.get(i) / counts.get(i));
        }
        return res;
    }

    public void dfs(TreeNode root, int level, List<Integer> counts, List<Double> sums) {
        if(root == null) return;
        if(level < sums.size()) {
            sums.set(level, sums.get(level) + root.val);
            counts.set(level, counts.get(level) + 1);
        } else {
            sums.add(root.val * 1.0);
            counts.add(1);
        }
        dfs(root.left, level+1, counts, sums);
        dfs(root.right, level+1, counts, sums);
    }
}