Tree Traversals iterative

tree_traversals_iterative.py

# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Stack(list):
    def push(self, node):
        self.append(node)
    
    def peek(self):
        return self[-1]


def iterativeDfsPreOrder(root):
    stack = Stack()
    
    stack.push(root)
    
    
    # Pre Order Traversal
    # Same as BFS, just with a stack instead of a Queue
    print ("\n\nPreOrder Root -> Left -> Right")
    while stack:
        current = stack.pop()
        print(current.val, end=' -> ')
        
        # In Reverse order
        if current.right:
            stack.push(current.right)
            
        if current.left:
            stack.push(current.left)
    
def iterativeDfsPostOrder(root):
    # Post Order Traversal
    print("\n\nPostOrder Left -> Right -> Root")
    stack = Stack()
    
    current = root
    
    while True:
        while current is not None:
            if current.right:
                stack.push(current.right)
            stack.push(current)
            current = current.left
        
        # current == None here in all cases
        current = stack.pop()
        
        if stack and stack.peek() == current.right:
            right_child = stack.pop()
            stack.push(current)
            current = right_child
        else:
            print(current.val, end = " -> ")
            # This ensures that the next interation picks up a new element from stack
            current = None
        
        
        if not stack and current is None:
            break

def iterativeDfsInOrder(root):
    # In Order Traversal
    print("\n\nInOrder Left -> Root -> Right")
    stack = Stack()
    
    current = root
    
    while True:
        # print([x.val for x in stack])
        while current is not None:
            if current.right:
                stack.push(current.right)
            stack.push(current)
            current = current.left
        # print([x.val for x in stack])
        
        # current == None here in all cases
        current = stack.pop()
        print(current.val, end=" -> ")
        
        if stack and stack.peek() == current.right:            
            right_child = stack.pop()
            current = right_child
        else:
            # This ensures that the next interation picks up a new element from stack
            current = None
        
        
        if not stack and current is None:
            break