Average of Levels in Binary Tree

Patrick Leaton
Published: April 3, 2021, 4:04 a.m.
Problem Description

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:

  1. The range of node's value is in the range of 32-bit signed integer.

 

The description was taken from https://leetcode.com/problems/average-of-levels-in-binary-tree/.

Problem Solution

#O(N) Time, O(M) Space, where M is Max Nodes
class Solution:
    def averageOfLevels(self, root: TreeNode) -> List[float]:
        if root is None:
            return []
        queue = collections.deque()
        queue.append(root)
        result = []
        while queue:
            level_sum = 0
            level_count = len(queue)
            for _ in range(level_count):
                node = queue.popleft()
                level_sum += node.val
                if node.left is not None:
                    queue.append(node.left)
                if node.right is not None:
                    queue.append(node.right)
            result.append(level_sum / level_count)
        return result

Problem Explanation

If we're needing to do a level order traversal, a Breadth-First Search is typically our go-to solution.  

During our BFS, we can have a level sum and level node count which we will reset before each level.  As we traverse each level, we will add the node value to the level sum.  Once we have finished a level traversal, we will append the level sum divided by the level count to the output. 

Let's start by checking that the root isn't null, if it is, we have no averages to return so we will return an empty list.

        if root is None:
            return []

 

Otherwise, we will initialize our queue, append the root to the queue, then initialize our result list.

        queue = collections.deque()
        queue.append(root)
        result = []

 

Now we can start our BFS that will continue while there are still nodes in the queue to visit.

        while queue:

 

At the beginning of each level iteration, we will reset our level_sum to be zero and also set the count of nodes within the level to be the length of nodes in the queue.

            level_sum = 0
            level_count = len(queue)

 

We will then iterate through each node in the queue.

            for _ in range(level_count):

 

During each node iteration, we will pop the node from the queue and add the node's value to the level sum.

                node = queue.popleft()
                level_sum += node.val

 

Afterward, if the node has children then we will push those onto the queue.

                if node.left is not None:
                    queue.append(node.left)
                if node.right is not None:
                    queue.append(node.right)

 

Once we have traversed each node in a level and have summed their values, we will take the average by dividing the level sum by the count of nodes in the level.

            result.append(level_sum / level_count)

 

Once we have traversed each node and have appended the average of each level to the result array, we will return the result.

        return result