算法题/树的BFS：二叉树每层的均值

题目

给定一棵二叉树，填充一个数组来表示它所有层次的平均值。 相似问题：找到二叉树的每一层节点中的最大值。

解答

Java

import java.util.*;

class TreeNode {
  int val;
  TreeNode left;
  TreeNode right;

  TreeNode(int x) {
    val = x;
  }
};

class LevelAverage {
  public static List<Double> findLevelAverages(TreeNode root) {
    List<Double> result = new ArrayList<>();
    if (root == null)
      return result;

    Queue<TreeNode> queue = new LinkedList<>();
    queue.offer(root);
    while (!queue.isEmpty()) {
      int levelSize = queue.size();
      double levelSum = 0;//将currentLevel换成levlSum
      for (int i = 0; i < levelSize; i++) {
        TreeNode currentNode = queue.poll();
        //添加当前节点的值到levelSum中
        levelSum += currentNode.val;
        //添加当前层的子节点
        if (currentNode.left != null)
          queue.offer(currentNode.left);
        if (currentNode.right != null)
          queue.offer(currentNode.right);
      }
      //添加当前层的值的平均值到结果数组中
      result.add(levelSum / levelSize);
    }

    return result;
  }

  public static void main(String[] args) {
    TreeNode root = new TreeNode(12);
    root.left = new TreeNode(7);
    root.right = new TreeNode(1);
    root.left.left = new TreeNode(9);
    root.left.right = new TreeNode(2);
    root.right.left = new TreeNode(10);
    root.right.right = new TreeNode(5);
    List<Double> result = LevelAverage.findLevelAverages(root);
    System.out.print("Level averages are: " + result);
  }
}

JavaScript

class TreeNode {
  constructor(val) {
    this.val = val;
    this.left = null;
    this.right = null;
  }
}

class LevelAverage {
  static findLevelAverages(root) {
    const result = [];
    if (root === null)
      return result;

    const queue = [];
    queue.push(root);
    while (queue.length > 0) {
      const levelSize = queue.length;
      let levelSum = 0;
      for (let i = 0; i < levelSize; i++) {
        const currentNode = queue.shift();
        levelSum += currentNode.val;
        if (currentNode.left !== null)
          queue.push(currentNode.left);
        if (currentNode.right !== null)
          queue.push(currentNode.right);
      }
      result.push(levelSum / levelSize);
    }

    return result;
  }
}

const root = new TreeNode(12);
root.left = new TreeNode(7);
root.right = new TreeNode(1);
root.left.left = new TreeNode(9);
root.left.right = new TreeNode(2);
root.right.left = new TreeNode(10);
root.right.right = new TreeNode(5);
const result = LevelAverage.findLevelAverages(root);
console.log("Level averages are: " + result);

Python

from collections import deque

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

class LevelAverage:
    def findLevelAverages(root):
        result = []
        if root is None:
            return result
        
        queue = deque()
        queue.append(root)
        while queue:
            level_size = len(queue)
            level_sum = 0
            for _ in range(level_size):
                current_node = queue.popleft()
                level_sum += current_node.val
                if current_node.left:
                    queue.append(current_node.left)
                if current_node.right:
                    queue.append(current_node.right)
            result.append(level_sum / level_size)
        
        return result

root = TreeNode(12)
root.left = TreeNode(7)
root.right = TreeNode(1)
root.left.left = TreeNode(9)
root.left.right = TreeNode(2)
root.right.left = TreeNode(10)
root.right.right = TreeNode(5)
result = LevelAverage.findLevelAverages(root)
print("Level averages are:", result)

Go

package main

import (
	"container/list"
	"fmt"
)

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

type LevelAverage struct{}

func (la *LevelAverage) findLevelAverages(root *TreeNode) []float64 {
	result := []float64{}
	if root == nil {
		return result
	}

	queue := list.New()
	queue.PushBack(root)
	for queue.Len() > 0 {
		levelSize := queue.Len()
		levelSum := 0.0
		for i := 0; i < levelSize; i++ {
			currentNode := queue.Remove(queue.Front()).(*TreeNode)
			levelSum += float64(currentNode.Val)
			if currentNode.Left != nil {
				queue.PushBack(currentNode.Left)
			}
			if currentNode.Right != nil {
				queue.PushBack(currentNode.Right)
			}
		}
		result = append(result, levelSum/float64(levelSize))
	}

	return result
}

func main() {
	root := &TreeNode{Val: 12}
	root.Left = &TreeNode{Val: 7}
	root.Right = &TreeNode{Val: 1}
	root.Left.Left = &TreeNode{Val: 9}
	root.Left.Right = &TreeNode{Val: 2}
	root.Right.Left = &TreeNode{Val: 10}
	root.Right.Right = &TreeNode{Val: 5}

	la := &LevelAverage{}
	result := la.findLevelAverages(root)
	fmt.Println("Level averages are:", result)
}