leetcode 637.二叉树的层平均值 Java

题目链接

https://leetcode-cn.com/problems/average-of-levels-in-binary-tree/

描述

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

示例

示例1

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

代码

BFS

/**
 * 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> list = new LinkedList<>();
        if (root == null){
            return list;
        }

        Deque<TreeNode> queue = new LinkedList<>();
        queue.add(root);

        while (!queue.isEmpty()){
            int size = queue.size();
            double sum = 0;
            for (int i = size; i > 0; i--){
                TreeNode node = queue.poll();
                sum += node.val;

                if (node.left != null) {
                    queue.offer(node.left);
                }
                if (node.right != null) {
                    queue.offer(node.right);
                }
            }

            list.add(sum / size);
        }

        return list;
    }
}

DFS

BFS的解法比较容易理解，然而这个题目DFS也是可以解决的。

用两个集合（不知道节点数，故不用数组），一个存放每一层的和，一个存放每一层的数目，用一个值作为区分节点属于哪个位置。

需要注意的也就是避免下标越界，dfs的时候加上判断就可以了。

class Solution {
    public List<Double> averageOfLevels(TreeNode root) {
        List<Double> res = new ArrayList<>();
        if (root == null) {
            return res;
        }

        //保存每一层的节点数和
        List<Double> sums = new ArrayList<>();
        //保存每一层节点的数目
        List<Integer> nums = new ArrayList<>();
        dfs(root, 0, sums, nums);

        for (int i = 0; i < sums.size(); i++){
            res.add(sums.get(i) / nums.get(i));
        }

        return res;
    }

    /**
    *   root：正在遍历的节点
    *   height：用来区分节点应该放置在何处，是sums第几个下标的值需要更新，nums第几个下标的值需要加一
    */
    private void dfs(TreeNode root, int height, List<Double> sums, List<Integer> nums){
        if (height < sums.size()){
            sums.set(height, sums.get(height) + root.val);
            nums.set(height, nums.get(height) + 1);
        } else {
            sums.add(1.0 * root.val);
            nums.add(1);
        }

        if (root.left != null) {
            dfs(root.left, height + 1, sums, nums);
        }

        if (root.right != null) {
            dfs(root.right, height + 1, sums, nums);
        }
    }
}