LeetCode 637: Average of Levels in Binary Tree

Problem Restatement

We are given the root of a binary tree.

We need to return an array where each element is the average value of all nodes on one level of the tree.

The first value is the average of the root level.

The second value is the average of the next level.

The result continues from top to bottom.

Answers within 10^-5 of the actual answer are accepted. The tree has at least one node, and node values may be large signed integers.

Input and Output

Example function shape:

def averageOfLevels(root: Optional[TreeNode]) -> list[float]:
    ...

Examples

Example 1:

root = [3, 9, 20, None, None, 15, 7]

The tree is:

        3
       / \
      9   20
          / \
         15  7

Level averages:

Output:

[3.0, 14.5, 11.0]

Example 2:

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

The output is also:

[3.0, 14.5, 11.0]

First Thought: Traverse the Whole Tree

We need to visit every node because every value contributes to exactly one level average.

The question is how to group nodes by level.

A depth-first search can work if we store:

level_sum[level]
level_count[level]

But breadth-first search fits the problem more directly because it naturally processes the tree one level at a time.

Key Insight

Use a queue.

At the start of each loop, the queue contains exactly the nodes of the current level.

So we can:

1. Read the current queue size.
2. Pop exactly that many nodes.
3. Sum their values.
4. Add their children to the queue.
5. Divide the sum by the number of nodes on that level.

This gives one average per loop iteration.

Algorithm

1. Create a queue containing the root node.
2. Create an empty result list.
3. While the queue is not empty:
   1. Store level_size = len(queue).
   2. Set level_sum = 0.
   3. Repeat level_size times:
      1. Pop one node.
      2. Add its value to level_sum.
      3. Push its left child if it exists.
      4. Push its right child if it exists.
   4. Append level_sum / level_size to the result.
4. Return the result.

Implementation

from collections import deque
from typing import Optional

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

class Solution:
    def averageOfLevels(self, root: Optional[TreeNode]) -> list[float]:
        queue = deque([root])
        answer = []

        while queue:
            level_size = len(queue)
            level_sum = 0

            for _ in range(level_size):
                node = queue.popleft()
                level_sum += node.val

                if node.left:
                    queue.append(node.left)

                if node.right:
                    queue.append(node.right)

            answer.append(level_sum / level_size)

        return answer

Code Explanation

We initialize the queue with the root:

queue = deque([root])

Since the problem guarantees at least one node, root is valid.

At the beginning of each loop:

level_size = len(queue)

This captures the number of nodes currently on this level.

Then we process exactly that many nodes:

for _ in range(level_size):

This is important. Children added during this loop belong to the next level, so they must not be included in the current level’s average.

Each node contributes to the current sum:

level_sum += node.val

Then we push its children:

if node.left:
    queue.append(node.left)

if node.right:
    queue.append(node.right)

After the level is complete, we compute the average:

answer.append(level_sum / level_size)

Python / returns a floating-point result, which matches the required output type.

Correctness

At the start of the first loop, the queue contains only the root, which is exactly level 0.

Assume that at the start of some loop, the queue contains exactly all nodes on level d, and no other nodes.

The algorithm stores the queue length as level_size, then pops exactly level_size nodes. Therefore, it processes every node on level d exactly once.

For each processed node, the algorithm adds its children to the queue. All children of level d nodes are exactly the nodes on level d + 1. Since the loop only processes the original level_size nodes, these children remain in the queue for the next iteration.

Thus, by induction, every loop processes exactly one level.

During each loop, level_sum is the sum of all node values on that level, and level_size is the number of nodes on that level. Therefore, level_sum / level_size is exactly the average for that level.

The algorithm appends these averages from top to bottom, so the returned list is correct.

Complexity

Let n be the number of nodes in the tree.

Here, w is the maximum width of the tree.

In the worst case, w can be O(n).

Alternative Solution: DFS With Sums and Counts

We can also use depth-first search and store the sum and count for each depth.

from typing import Optional

class Solution:
    def averageOfLevels(self, root: Optional[TreeNode]) -> list[float]:
        sums = []
        counts = []

        def dfs(node: Optional[TreeNode], depth: int) -> None:
            if node is None:
                return

            if depth == len(sums):
                sums.append(0)
                counts.append(0)

            sums[depth] += node.val
            counts[depth] += 1

            dfs(node.left, depth + 1)
            dfs(node.right, depth + 1)

        dfs(root, 0)

        return [sums[i] / counts[i] for i in range(len(sums))]

This also visits every node once.

Here, h is the tree height, and L is the number of levels.

Testing

def run_tests():
    s = Solution()

    root = TreeNode(
        3,
        TreeNode(9),
        TreeNode(20, TreeNode(15), TreeNode(7)),
    )
    assert s.averageOfLevels(root) == [3.0, 14.5, 11.0]

    root = TreeNode(1)
    assert s.averageOfLevels(root) == [1.0]

    root = TreeNode(
        -1,
        TreeNode(-2),
        TreeNode(-3),
    )
    assert s.averageOfLevels(root) == [-1.0, -2.5]

    root = TreeNode(
        5,
        TreeNode(14, TreeNode(1), None),
        None,
    )
    assert s.averageOfLevels(root) == [5.0, 14.0, 1.0]

    print("all tests passed")

run_tests()

Test meaning: