Binary Tree Level Order Traversal in Python

user_5712621 0
Viewed 2618

Given the following tree:

enter image description here

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
tree breadth-first-search data-structures python
3 Answers

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]]
user_106560540

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
user_57434060

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 .

user_126536980
Related