Binary Tree Level Order Traversal in Python

Question

Given the following tree:

I should return the level order traversal of the tree from left to right: so the above example would output a list of lists :

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

I wrote the following code, idea is storing the node value and its depth recursively in a Queue then iterate the Queue tuples and put the in intermediate list O if no more node of same depth append O to output and empty O and so on. However my code Timeout any help?

import queue

class Solution:
    def levelOrder(self, root: TreeNode) -> List[List[int]]:
        def helper(root,res,level):
            if not root:
                return res
            l=level+1
            res.put((root.val,level))
            helper(root.left,res,l)
            helper(root.right,res,l)

        res=queue.Queue()
        helper(root,res,0)

        d=1
        output=[]
        node,depth=res.get()
        output.append([node])
        while res:
            o=[]
            node,depth=res.get()
            while d ==depth:
                o.append(node)
                node,depth=res.get()
            else:
                d+=1
            output.append(o)    
            
        return output

Answer 1

Thanks for the answer from san. The solution helps me solve problem 107 of leetcode. I revised a little bit based on my understanding. There are two ways to solve this question but I prefer to use BFS fashion. This revised version compresses both value and level in the queue. It provides more flexibility comparing with only print the value of each node in the tree, as other tutorials provided.

# 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 levelOrderBottom(self, root: TreeNode) -> List[List[int]]: 
        
        if root is None:
            return None
        
        level = 1
        q = []
        q.append((root,level)) # push both node and level.
        save = []
        
        while q:
            cur = q.pop(0)
            cur_node = cur[0]
            cur_level = cur[1]
            # print(cur_node.val)
            # print(cur_level)  
            save.append([cur_node.val, cur_level])
            if cur_node.left:
                q.append((cur_node.left, cur_level+1))
            if cur_node.right:
                q.append((cur_node.right, cur_level+1))
                
        print(save) # once print, you will have the idea about how to reorgnized the required output.
        level = 1
        output = []
        temp = []
        for i in range(len(save)):
            cur = save[i]
            #print(cur)
            if cur[1] == level:
                temp.append(cur[0])    
            if cur[1] != level:
                output.insert(0, temp)
                temp = []
                temp.append(cur[0])
                level = level + 1
            if i == len(save)-1:
                    output.insert(0, temp)
            
        return output

Answer 2

You can use the below logic , to print out nodes in BFS . If you also need , you modify the method to return a list as well .

def levelOrder(root):
qroot = []
print(root.info,end=' ')
if root.left:
    qroot.append(root.left)
if root.right:
    qroot.append(root.right)

while(qroot):
    tmp = qroot[0]
    if tmp.left:
        qroot.append(tmp.left)
    if tmp.right:
        qroot.append(tmp.right)

    print (tmp.info,end=' ')
    qroot.pop(0)

return

Instead of printing the values , you can directly keep appending the value to a new list and return the same .

Answer 3

Here is my code for the Breadth First Search (BFS) Iterative Implementation with nodes in each level in the final output is under a single list:

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

class Solution:
    def BFS(self, root) -> int:
        level=1
        current=(root, level)
        s=set()
        result=[]
        Q = [current]
        while Q:
            current=Q.pop()
            level=current[1]
            if current[0] not in s:
                result.append([current[0].val, level])
                s.add(current[0])
            if current[0].left:
                Q.insert(0,(current[0].left, level+1))
            if current[0].right:
                Q.insert(0,(current[0].right, level+1))
        output=[]
        temp=[]
        level=1
        for val in result:
            if val[1]==level:
                temp.append(val[0])
            elif val[1] > level:
                output.append(temp)
                temp=[val[0]]
                level+=1
        output.append(temp)
        return output

Testing:

n1=TreeNode(3)
n2=TreeNode(9)
n3=TreeNode(20)
n4=TreeNode(6)
n5=TreeNode(15)
n6=TreeNode(7)

n1.left=n2
n1.right=n3
n2.left=n4
n3.left=n5
n3.right=n6

sol1=Solution()
print(sol1.BFS(n1))
[[3], [9, 20], [6, 15, 7]]