题目:
给定一个非空二叉树, 返回一个由每层节点平均值组成的数组.
示例 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);
}
}