LeetCode-in-Go

637. Average of Levels in Binary Tree

Easy

Given the root of a binary tree, return the average value of the nodes on each level in the form of an array. Answers within 10-5 of the actual answer will be accepted.

Example 1:

Input: root = [3,9,20,null,null,15,7]

Output: [3.00000,14.50000,11.00000] 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].

Example 2:

Input: root = [3,9,20,15,7]

Output: [3.00000,14.50000,11.00000]

Constraints:

Solution

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func averageOfLevels(root *TreeNode) []float64 {
	levelSums := make(map[int][]float64)
	helper(root, levelSums, 0)
	result := make([]float64, len(levelSums))
	for level, pair := range levelSums {
		result[level] = pair[1] / pair[0]
	}
	return result
}

func helper(root *TreeNode, levelSums map[int][]float64, level int) {
	if root == nil {
		return
	}
	if _, exists := levelSums[level]; !exists {
		levelSums[level] = []float64{0, 0}
	}
	pair := levelSums[level]
	pair[0]++                    // count
	pair[1] += float64(root.Val) // sum
	levelSums[level] = pair
	helper(root.Left, levelSums, level+1)
	helper(root.Right, levelSums, level+1)
}