Average of Levels in Binary Tree

Difficulty: Easy

Understanding The Problem

Problem Description

Given a binary tree, write a program to return the average value of the nodes on each level in the form of an array.

Problem Note

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

Example

Input:
    4
   / \
  8  11
    /  \
   13   7
Output: [4, 9.5, 10]
Explanation: The average value of nodes on level 0 is 4, on level 1 is 9.5, and on level 2 is 10. Hence return [4, 9.5, 10].

The node structure passed to the function is:

struct TreeNode{
    int val
    TreeNode left
    TreeNode right
}

You can try practising this problem here.

Solution

You must have read about the Level Order Traversal in a tree that is used to traverse in a binary tree. The level order traversal is similar to BFS traversal where the starting point is the root node. We can use this traversal technique here.

How can we use Level Order Traversal Technique to solve the problem?

We will use a Queue that will store the nodes of each level. We will enqueue all the nodes of the current level while storing their sum in a variable. The total number of nodes in the current level will be equal to the queue size. Using them we can find the average of the current level and store it in the result array. Afterwards, we will dequeue all the nodes from Queue and enqueue the children of the dequeued nodes which will represent the next level of the tree.

Solution Step

Create a Queue and push the root node in it. Create an empty result array and append the root node value in it.
Find the length of the Queue and store it in a currLevelNoOfNodes variable

Dequeue all the nodes of Queue and enqueue all the children of the dequeued nodes
During enqueue, keep a sum variable that will store the sum of all the nodes of the current level.
For each such step, append the value (sum/currLevelNoOfNodes) in the result array.

3. Repeat step 2 until the Queue becomes empty.

Pseudo Code

float[] averageOfLevels(TreeNode root){
    float[] result
    Queue q
    q.push(root)
    result.append(root.val)
    while(q is not empty){
        sum = 0
        prevLevelNoOfNodes = q.size()
        currLevelNoOfNodes = 0
        while(prevLevelNoOfNodes > 0){
            TreeNode node = q.pop()
            prevLevelNoOfNodes = prevLevelNoOfNodes - 1
            if(node.left){
                q.push(node.left)
                sum = sum + node.left.val
                currLevelNoOfNodes = currLevelNoOfNodes + 1
            }
            if(node.right){
                q.push(node.right)
                sum = sum + node.right.val
                currLevelNoOfNodes = currLevelNoOfNodes + 1
            }
        }
        if(currLevelNoOfNodes > 0){
            result.append(sum/currLevelNoOfNodes)
        }
    }
    return result
}

Complexity Analysis

Time Complexity: O(n)

Space Complexity: O(n)

Critical Ideas To Think

How did the Level Order Traversal help to find the average of each level? Answer : Level order traversal helped to traverse level by level thereby giving average values of nodes at each level.
In pseudocode, we appended in result array only when currLevelNoOfNodes > 0 ? Answer : This is because if the currLevelNoOfNodes is 0, that means there were no nodes at that level and thereby it will not contribute to the result array as all the levels of the tree been traversed by then.
How come the prevLevelNoOfNodes are equal to the queue size?
Why did we append the root value in the result first? Answer : The number of nodes at the root level is 1 and thus the average will basically be the value of the root itself at the first level.