637. Average of Levels in Binary Tree #

Problem #

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

Example 1:


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.

Problem Summary #

Traverse a tree level by level from top to bottom, and calculate the average value of each level.

Solution Approach #

This can be implemented with a queue.
Problems 102 and 107 are both level-order traversal problems.

Code #


package leetcode

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func averageOfLevels(root *TreeNode) []float64 {
	if root == nil {
		return []float64{0}
	}
	queue := []*TreeNode{}
	queue = append(queue, root)
	curNum, nextLevelNum, res, count, sum := 1, 0, []float64{}, 1, 0
	for len(queue) != 0 {
		if curNum > 0 {
			node := queue[0]
			if node.Left != nil {
				queue = append(queue, node.Left)
				nextLevelNum++
			}
			if node.Right != nil {
				queue = append(queue, node.Right)
				nextLevelNum++
			}
			curNum--
			sum += node.Val
			queue = queue[1:]
		}
		if curNum == 0 {
			res = append(res, float64(sum)/float64(count))
			curNum, count, nextLevelNum, sum = nextLevelNum, nextLevelNum, 0, 0
		}
	}
	return res
}