二叉树的层平均值

Average of Levels in Binary Tree ⭐️

LeetCode每日一题 2020.09.12

Given a non-empty binary tree, return the average value of the nodes on each level in the form of an array.

Example:

Input:
    3
   / \
  9  20
    /  \
   15   7
Output: [3, 14.5, 11]
Explanation:
The average value of nodes on level 0 is 3,  on level 1 is 14.5, and on level 2 is 11. Hence return [3, 14.5, 11].

note: The range of node’s value is in the range of 32-bit signed integer.

题目大意

按照题目大意需要计算二叉树每层所有节点的平均值, 必然需要遍历整个二叉树节点, 所以自然想到DFS(深度优先检索)和BFS(广度优先检索)两种方式。

解题思路

方式一: DFS

使用深度优先检索算法计算二叉树的层平均值, 首先定义两个数组, 第一个数组(counts)用于存储二叉树的每一层的节点数, 第二个数组(sums)用于存储二叉树的每一层的节点值之和。遍历完整个二叉树之后，第 i 层的平均值即为 sums[i]/counts[i].

方式二: BFS

使用广度优先检索算法计算二叉树的层平均值. 用层序遍历的方法, 维护一个队列去遍历节点. 用 for 循环控制一层的节点逐个出列, 节点值累加求和. 节点出列的同时, 下一层的子节点加入队列, 在 for 循环结束时，队列中就全是下一层的节点. 此时当前层的求和也好了, 除以当前层的节点个数, 就是当前层的平均值, 加入结果数组. 接着处理下一层的节点, 重复以上步骤, 就构建好了结果数组.

广度优先检索算法通过使用队列存储待访问节点, 只要保证在每一轮遍历时, 当前队列中的节点是同一层的所有节.

第一步, 将根节点加入队列;
每一轮遍历, 将队列中的所有节点取出, 计算节点的数量以及节点值之和, 同时计算节点的平均值, 然后将节点的全部非空子节点重新加入队列, 直到队列为空, 遍历结束.

代码

DFS

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */

type data struct{sum, count int}

func averageOfLevels(root *TreeNode) []float64 {
    levelsData := []data{}
    var dfs func(node *TreeNode, level int)
    dfs = func(node *TreeNode, level int) {
        if node == nil {
            return
        }
        if level < len(levelsData) {
            levelsData[level].sum += node.Val
            levelsData[level].count++
        } else {
            levelsData = append(levelsData, data{node.Val, 1})
        }
        dfs(node.Left, level+1)
        dfs(node.Right, level+1)
    }

    dfs(root, 0)
    averages := make([]float64, len(levelsData))
	for i, d := range levelsData {
		averages[i] = float64(d.sum) / float64(d.count)
	}
	return averages
}

BFS

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */

func averageOfLevels(root *TreeNode) (averages []float64) {
	nextLevel := []*TreeNode{root}
	for len(nextLevel) > 0 {
		sum := 0
		curLevel := nextLevel
		nextLevel = nil
		for _, node := range curLevel {
			sum += node.Val
			if node.Left != nil {
				nextLevel = append(nextLevel, node.Left)
			}
			if node.Right != nil {
				nextLevel = append(nextLevel, node.Right)
			}
		}
		averages = append(averages, float64(sum)/float64(len(curLevel)))
	}
	return
}