add test case and clean code for BFS python

Author: kumailn

Python/Binary_Tree_Level_Order_Traversal.py

@@ -7,19 +7,26 @@ def levelOrder(root):
         #Return empty list if the root node is null
         if not root: return []
         #vaiable to hold final result, queue of nodes and current level of tree
-        r, q, c = [[]], [root], []
+        order, queue, nextQueue = [], [root], []
         #While the queue is not empty
-        while q:
+        while queue:
             #For every node in the queue
-            for i in range(len(q)):
-                #Append the node to the last list in the result list
-                r[-1] += [q[i].val]
+            for node in range(len(queue)):
                 #If that node has a left child, append to current list, then same for right child
-                c = c + [q[i].left] if q[i].left else c
-                c = c + [q[i].right] if q[i].right else c
-            #Let queue be the current traversal level nodes, and current level resets to empty list
-            q, c = c, []
-            #Append empty list to result (for next level) if queue is non empty
-            r = r + [[]] if q else r
-        return r
-        
+                if queue[node].left: nextQueue += [queue[node].left]
+                if queue[node].right: nextQueue += [queue[node].right]
+            #Now we append the current queue to the order (the lambda is just there to turn the TreeNode objects to their values)
+            order += [list(map(lambda x: x.val, queue))]
+            queue, nextQueue = nextQueue, []
+        return order
+        
+def main():
+    tree = TreeNode(6)
+    tree.left = TreeNode(3)
+    tree.right = TreeNode(12)
+    tree.left.right = TreeNode(4)
+    tree.right.right = TreeNode(14)
+
+    print(levelOrder(tree))
+
+main()