Search code examples
pythonpython-3.xdata-structurestreebinary-tree

Python nested loop not breaking out for level order tree traversal


I am solving a leetcode problem: https://leetcode.com/problems/binary-tree-level-order-traversal/ I am using a list as a queue to solve this problem. I am using two loops to solve this question, the inner loop breaks each time a level of the tree is traversed

Here's a brief overview of the code:

  • The q list is used as a queue to store nodes.
  • The root node is initially appended to the queue.
  • A marker (-1) is appended to the queue to indicate the end of the first level.
  • The algorithm continues until the queue is empty.
  • Inside the outer loop, there is an inner loop to process nodes at the current level.
  • Nodes are dequeued one by one, and their values are appended to the cur list.
  • If the dequeued node is the marker (-1), it means the current level is complete, and cur is appended to the result (res), and a new marker is added to the queue for the next level.
  • If the dequeued node is a valid tree node, its left and right children are added to the queue.
  • The process continues until the entire tree is traversed.

The problem is the program wont quit, even after getting into the break statements. In the logs it prints out that it is "breaking outer loop" and still continues where it was supposed to end the program. Moreover, it is getting a detecting a node with value 1, which was never there in the input (See the code, input and logs below)

This is the input to the problem enter image description here

Code I am using

def levelOrder(self, root: Optional[TreeNode]) -> List[List[int]]:
        q = []
        q.append(root)
        q.append(-1)
        res = []
        while len(q) != 0:
            cur = []
            while len(q) != 0:
                node = q.pop(0)
                print(q, node)
                if node == -1:
                    if len(q) == 0:
                        print("End")
                        break
                    res.append(cur)
                    q.append(-1)
                    break
                cur.append(node.val)
                if node.left != None:
                    q.append(node.left)
                if node.right != None:
                    q.append(node.right)
            if len(q) == 0:
                print("break outer loop")
                break
        return res

Logs of the code

[-1] TreeNode{val: 3, left: TreeNode{val: 9, left: None, right: None}, right: TreeNode{val: 20, left: TreeNode{val: 15, left: None, right: None}, right: TreeNode{val: 7, left: None, right: None}}}
[TreeNode{val: 9, left: None, right: None}, TreeNode{val: 20, left: TreeNode{val: 15, left: None, right: None}, right: TreeNode{val: 7, left: None, right: None}}] -1
[TreeNode{val: 20, left: TreeNode{val: 15, left: None, right: None}, right: TreeNode{val: 7, left: None, right: None}}, -1] TreeNode{val: 9, left: None, right: None}
[-1] TreeNode{val: 20, left: TreeNode{val: 15, left: None, right: None}, right: TreeNode{val: 7, left: None, right: None}}
[TreeNode{val: 15, left: None, right: None}, TreeNode{val: 7, left: None, right: None}] -1
[TreeNode{val: 7, left: None, right: None}, -1] TreeNode{val: 15, left: None, right: None}
[-1] TreeNode{val: 7, left: None, right: None}
[] -1
End
break outer loop <- The program was supposed to end here.
[-1] TreeNode{val: 1, left: None, right: None} <- No clue from where it is getting a node with value 1
[] -1
End
break outer loop
[-1] None

I was expecting the program to end when the outer loop breaks off, even if I set a flag in the inner loop, to end flag and break off, or even when I directly return the answer in the inner loop, the program still continues the execution..


Solution

  • The output is expected. Your function does quit, but you are looking at the output of mulitple test cases that are run one after the other. You could easily see that if you would put a print statement as the first statement in your function, like:

    print("----start of function----")
    

    There are two problems that remain in your code:

    • It does not deal well with an empty tree. If root is None your code will fail. Just add an if that tests for this condition and then returns an empty list.

    • It does not add the last level to the output. Where you have print("End"), you didn't add cur to res, yet there are still values in cur. So move this if block below the statement that does that.

    With those two changes your code will work.

    class Solution:
        def levelOrder(self, root: Optional[TreeNode]) -> List[List[int]]:
            if not root:  # deal with boundary case
                return []
            q = []
            q.append(root)
            q.append(-1)
            res = []
            while len(q) != 0:
                cur = []
                while len(q) != 0:
                    node = q.pop(0)
                    if node == -1:
                        res.append(cur) 
                        if len(q) == 0:  # Only exit AFTER adding cur
                            break
                        q.append(-1)
                        break
                    cur.append(node.val)
                    if node.left != None:
                        q.append(node.left)
                    if node.right != None:
                        q.append(node.right)
                if len(q) == 0:
                    break
            return res