Average of Levels in Binary Tree

Problem Id: 637 Difficulty: Easy Tag: Tree


Intuition

Use BFS to solve this problem. The only one difference is to add an additional variable in the BFS queue which is the depth of the current node.

In BFS, we will traversal each level's node at the same time, so we could get the sum of the current level. Once the level changes, we could add the average value to the result.

Solution


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
from collections import deque


class Solution:
    def averageOfLevels(self, root: TreeNode) -> List[float]:
        queue = deque([])
        queue.append([root, 0])
        ans = []
        level_sum = 0.0
        level_count = 0
        level = 0
        while queue:
            node, depth = queue.popleft()
            if depth != level:
                ans.append(level_sum / level_count)
                level_sum, level_count, level = 0.0, 0, depth
            level_sum += node.val
            level_count += 1
            if node.left:
                queue.append([node.left, depth + 1])
            if node.right:
                queue.append([node.right, depth + 1])
        ans.append(level_sum / level_count)
        return ans