Leetcode – Average of Levels in Binary Tree Solution in Python

Spread the love

Introduction

Binary trees are among the most fundamental and widely used data structures in computer science. In this article, we will tackle a popular binary tree problem from Leetcode named “Average of Levels in Binary Tree”. This problem is an excellent way to understand tree traversal techniques and learn how to solve problems that require aggregating information across different levels of a binary tree. We will analyze the problem statement, consider various solutions, and walk through Python code examples.

Problem Statement (Leetcode #637)

Given a non-empty binary tree, return the average value of the nodes on each level in the form of an array.

Example:

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 (9 + 20)/2 = 14.5,
  • And on level 2 is (15 + 7)/2 = 11.

Definition for a binary tree node:

In Python, a binary tree node can be defined using a class as follows:

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

Solution Approaches

1. Breadth-First Search (BFS) Approach

One of the most intuitive ways to solve this problem is by using Breadth-First Search (BFS). In BFS, we traverse the tree level by level, which aligns perfectly with the requirements of this problem.

  1. Initialize an empty list to store the average values at each level.
  2. Initialize a queue and enqueue the root node.
  3. While the queue is not empty, perform the following steps: a. Determine the number of nodes at the current level (i.e., the size of the queue). b. Initialize a variable to store the sum of node values at the current level. c. For each node at the current level, dequeue the node, add its value to the sum, and enqueue its non-null left and right children. d. Calculate the average value at the current level and append it to the list.
  4. Return the list of average values.
from collections import deque

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: TreeNode):
        if not root:
            return []
        
        averages = []
        queue = deque([root])
        
        while queue:
            level_length = len(queue)
            level_sum = 0
            for i in range(level_length):
                node = queue.popleft()
                level_sum += node.val
                if node.left:
                    queue.append(node.left)
                if node.right:
                    queue.append(node.right)
            averages.append(level_sum / level_length)
            
        return averages

Time Complexity:

  • O(n), where n is the number of nodes in the tree. Each node is enqueued and dequeued once.

Space Complexity:

  • O(m), where m is the maximum number of nodes in any level of the tree. This is the maximum size the queue can grow to.

2. Depth-First Search (DFS) Approach

An alternative approach is using Depth-First Search (DFS). Although less intuitive for this problem, DFS can still be used effectively.

  1. Initialize two lists: one to store the sum of values at each level, and another to store the number of nodes at each level.
  2. Perform a DFS traversal, keeping track of the current level. For each node, add its value to the corresponding sum and increment the node count for the level.
  3. After the traversal, calculate the average for each level.
  4. Return the list of average values.
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: TreeNode):
        def dfs(node, level):
            if not node:
                return
            
            if level < len(sums):
                sums[level] += node.val
                counts[level] += 1
            else:
                sums.append(node.val)
                counts.append(1)
                
            dfs(node.left, level + 1)
            dfs(node.right, level + 1)
        
        sums = []
        counts = []
        dfs(root, 0)
        
        return [s / c for s, c in zip(sums, counts)]

Time Complexity:

  • O(n), where n is the number of nodes in the tree. Each node is visited once.

Space Complexity:

  • O(h), where h is the height of the tree. This accounts for the maximum size of the recursive call stack.

Conclusion

In this article, we explored the “Average of Levels in Binary Tree” problem on Leetcode and discussed two different approaches in Python: Breadth-First Search and Depth-First Search. BFS is the more intuitive solution, as it naturally handles level-wise traversal. However, DFS is a viable alternative that demonstrates the flexibility of depth-first traversal techniques. Understanding and being proficient with both BFS and DFS is crucial for solving a wide range of tree-based problems effectively.

Leave a Reply