Average of Levels in Binary Tree Solution
Leetcode Solution: Understand and solve Leetcode problem Average of Levels in Binary Tree(637)
In this article, we’ll be solving the problem: Average of Levels in Binary Tree. This problem belongs to Tree Traversal and Recursion category.
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.
The solution to this problem will require a tree traversal and each of the node that we’ll encounter while traversing must know their own height in the tree. The height is important here because we’ve to return the average value of node for each level of the tree.
One more important point to look into is the formula for calculating the average using the previous average, and the new value:
NEW_AVERAGE = (
OLD_AVERAGE x PREVIOUS_NUMBER_OF_ELEMENTS + NEW_VALUE
)
/
(PREVIOUS_NUMBER_OF_ELEMENTS + 1)
I’ll keep the average per height into a map. The key for the map will be the height, and the value will be a tuple that will contain the current average and number of elements at that level. These values will be required to calculate the new average whenever we encounter a new node in the tree.
So, to solve this problem:
- I’ll traverse the tree using pre Order approach
- While traversing, I’ll keep the height for the current node available so that I can use the value of current height to access the map.
- Once a new node is found, I’ll get the average for the current height then I’ll update it and put it back in the map.
- This way once the traversal ends, the map will store average node values for all the heights.
- I’ll traverse this map and create a list of average values so that I can return it as Leetcode expects the result that way.
Here is the code for the solution whichis available here on GitHub as well
type TreeNode struct {
Val int
Left *TreeNode
Right *TreeNode
}
var levelAvgNodeCache map[int][2]float64
func averageOfLevels(root *TreeNode) []float64 {
levelAvgNodeCache = make(map[int][2]float64)
if root == nil {
return []float64{}
}
preOrder(root, 0)
var avgValues = make([]float64, len(levelAvgNodeCache))
// create a list from the map
for k, v := range levelAvgNodeCache {
avgValues[k] = v[0]
}
return avgValues
}
func preOrder(node *TreeNode, height int) {
values, ok := levelAvgNodeCache[height]
if !ok {
// if I encounter a node for a height for the first time
// I just create an entry for it
levelAvgNodeCache[height] = [2]float64{float64(node.Val), 1}
} else {
// calculating the new average here
newAvg := (float64(node.Val) + (values[0] * values[1])) / (values[1] + 1)
levelAvgNodeCache[height] = [2]float64{newAvg, values[1] + 1}
}
if node.Left != nil {
// traverse the left node with increased height
preOrder(node.Left, height + 1)
}
if node.Right != nil {
// traverse the right node with increased height
preOrder(node.Right, height + 1)
}
}
