Zig-Zag Level Order Traversal of BST/Binary-Tree [O(N)]

zig_zag_level_order.py

class TreeNode:
    def __init__(self, data, depth):
        self.data = data
        self.depth = depth
        self.left = None
        self.right = None
        
        
def BST_insert_helper(root, node):
    data = node.data
    node.depth += 1
    if data < root.data:
        if root.left is not None: 
            BST_insert_helper(root.left, node)
        else: 
            root.left = node
    else:
        if root.right is not None: 
            BST_insert_helper(root.right, node)
        else: 
            root.right = node
        
def BST_insert(root, data):
    node = TreeNode(data, 0)
    if root is None: 
        root = node
    else: 
        temp = root
        BST_insert_helper(temp, node)
    
def level_order(Q):
    res = []  # for each level, insert the nodes at the respective level.
    count = [False for x in range(1001)]
    while Q:
        node = Q.pop(0)
        if count[node.depth] is False:
            count[node.depth] = True
            temp = [node.data]
            res.append(temp)
        else: res[node.depth].append(node.data);
        
        if node.left is None and node.right is None: continue;
        elif node.left is None: Q.append(node.right);
        elif node.right is None: Q.append(node.left);
        else:
            Q.append(node.left)
            Q.append(node.right)
    
    return res

# FOR TESTING PURPOSE
def inorder(root):
    if root is not None:
        inorder(root.left)
        print(root.data, root.depth)
        inorder(root.right)

if __name__ == "__main__":
    t = int(input())
    while t > 0:
        t -= 1
        n = int(input())
        l = list(map(int, input().strip().split(' ')))
        ROOT = TreeNode(l[0], 0)
        for i in range(1, len(l)): 
            BST_insert(root=ROOT, data=l[i])
        # inorder(ROOT)
        result = level_order(Q=[ROOT])
        for i in range(0, len(result)):
            if i%2 == 1: 
                for j in range(0, len(result[i])): 
                    print(result[i][j], end=" ");
            else: 
                for j in range(len(result[i])-1, -1, -1): 
                    print(result[i][j], end=" ");
        if t > 0: print("")