[Leetcode] 637. Average of Levels in Binary Tree

https://leetcode.com/problems/average-of-levels-in-binary-tree/submissions/1330130110/?envType=study-plan-v2&envId=top-interview-150

initial

Just bfs traversal and storing the information of values and level(depth of tree) in a dictionary, where key would be the level and value would be a [total_sum,total_count_of_nodes_inthat_level]. I just stored total sum as value cuz I didnt think we need to store number of children (cuz of the formula 2i) but I didnt think of the case where the children can be null. If children are null, the node numbers for that level can be less than 2i.

in simple terms, number of children nodes must be known for each level to calculate the average.

from collections import deque
from typing import Optional, List

# 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]:
        dic = {}
        queue = deque()
        queue.append((root, 1))
        self.bfs(queue, dic)

        ans = [dic[i] / 2**(i-1) for i in dic.keys()]
        return ans

    def bfs(self, queue, dic):
        while queue:
            node, level = queue.popleft()
            
            val = node.val
            if level not in dic:
                dic[level] =val
            else:
                dic[level] += val
            
            if node.left:
                queue.append((node.left, level + 1))
            if node.right:
                queue.append((node.right, level + 1))

# Example usage
# Tree:
#       1
#      / \
#     2   3
#    / \   \
#   4   5   6

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)
root.right.right = TreeNode(6)

solution = Solution()
print(solution.averageOfLevels(root))  # Output: [1.0, 2.5, 5.0]

solution

class Solution:
    def averageOfLevels(self, root: Optional[TreeNode]) -> List[float]:
        if not root:
            return []

        # Use a dictionary to store the sum and count of nodes at each level
        dic = {}
        queue = deque([(root, 0)])  # Use level 0 for the root

        while queue:
            node, level = queue.popleft()

            if level not in dic:
                dic[level] = [0, 0]  # Initialize sum and count

            # Update sum and count for the current level
            dic[level][0] += node.val
            dic[level][1] += 1

            # Add child nodes to the queue with the incremented level
            if node.left:
                queue.append((node.left, level + 1))
            if node.right:
                queue.append((node.right, level + 1))

        # Calculate the average for each level
        ans = [dic[level][0] / dic[level][1] for level in dic]
        return ans
        

complexity

Time Complexity
The time complexity of the averageOfLevels function can be analyzed as follows:

Breadth-First Search (BFS) Traversal:

The BFS traversal visits each node exactly once. Since there are
𝑛
n nodes in the tree, this part of the algorithm takes
𝑂
(
𝑛
)
O(n) time.
Dictionary Updates:

Each node is processed, and the dictionary dic is updated with the node's value and count at the corresponding level. This update operation is
𝑂
(
1
)
O(1) per node, so the total time for updating the dictionary is
𝑂
(
𝑛
)
O(n).
Calculating Averages:

After the BFS traversal, the function calculates the average for each level by iterating over the dictionary dic. In the worst case, the tree could have
𝑛
n levels (in the case of a skewed tree). However, in a balanced tree, the number of levels
𝐿
L is
𝑂
(
log
⁡
𝑛
)
O(logn). Therefore, iterating over the dictionary takes
𝑂
(
𝐿
)
O(L) time, which is
𝑂
(
𝑛
)
O(n) in the worst case.
Combining these steps, the overall time complexity is:
𝑂
(
𝑛
)
O(n)

Space Complexity
The space complexity of the averageOfLevels function can be analyzed as follows:

Queue:

The queue is used to perform the BFS traversal. In the worst case, the maximum number of nodes in the queue at any point in time is the maximum number of nodes at any level in the tree. For a balanced binary tree, this is approximately
𝑛
/
2
n/2, leading to
𝑂
(
𝑛
)
O(n) space. For a skewed tree, it is
𝑂
(
𝑛
)
O(n).
Dictionary dic:

The dictionary stores sums and counts for each level. In the worst case, the tree has
𝑛
n levels, so the dictionary will store
𝑛
n entries. Each entry is a list of two integers, so the space taken by the dictionary is
𝑂
(
𝑛
)
O(n).
Combining these factors, the overall space complexity is:
𝑂
(
𝑛
)
O(n)

Summary
Time Complexity:
𝑂
(
𝑛
)
O(n)
Space Complexity:
𝑂
(
𝑛
)
O(n)
This ensures that the algorithm is efficient in terms of both time and space, making it suitable for processing binary trees of various sizes.