[Top Interview 150] 637. Average of Levels in Binary Tree

Solution)

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

        while queue:
            nodeSum = 0
            nodeCount = len(queue)

            for i in range(nodeCount):
                cur = queue.popleft()

                nodeSum += cur.val

                if cur.left: queue.append(cur.left)
                if cur.right: queue.append(cur.right)

            res.append(nodeSum/nodeCount)
        
        return res

 
Problem: LeetCode
 
Time Complexity: O(n)
Space Complexity: O(n)
 
*n = number of nodes

Explanation)

We'll use BFS to find nodes on each level using dequeue.
 

Initialization)

We'll have a dequeue (for the sake of simplicity, we'll call this dequeue a queue) that will save nodes on each level.
There'll also be an empty list that will save the average of the nodes.
 
The reason why we use a dequeue is because it's faster than a normal list. We need to pop the first element of dequeue when finding a node. We can use popleft() method for a dequeue, which is O(1) time complexity, while we have to use pop(0) method for a list, which is O(n) time compexity. 
 

 

def averageOfLevels(self, root: Optional[TreeNode]) -> List[float]:
        queue = deque([root])
        res = []

 

 

While Loop)

A while loop will be continued as long as there is a node in queue.
Having nodes in queue means that there are nodes that have not been calculated yet.
 
Inside the while loop, there will be another for loop that will iterate nodes on each level of a tree
 
The for loop will be looped as many times as the number of nodes on that level.
The loop will pop nodes of that level, sum up the values of those nodes, and append their children nodes into queue.
 
After all the nodes on each level are summed up, we'll calculate the average of the sum and append it to res.
 
If the while loop is done, meaning all the nodes in the tree are calculated, we'll return res.
 

 

def averageOfLevels(self, root: Optional[TreeNode]) -> List[float]:
        queue = deque([root])
        res = []

        while queue:
            nodeSum = 0
            nodeCount = len(queue)

            for i in range(nodeCount):
                cur = queue.popleft()

                nodeSum += cur.val

                if cur.left: queue.append(cur.left)
                if cur.right: queue.append(cur.right)

            res.append(nodeSum/nodeCount)
        
        return res

 

 
TMI: I'm not really used to working with BFS, so it was hard to get the structure of it. But once I got the slight sketch of it, everything was put in-place. Other than implementing a for loop, everything else was fine.