List methods for traversaling trees in various orders

tree_traversal.py

#!python3
from collections import deque


class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None


def preorder_iterative(root):
    if not root:
        return []
    stack = [root]
    res = []
    while stack:
        n = stack.pop()
        res.append(n.val)
        if n.right:
            stack.append(n.right)
        if n.left:
            stack.append(n.left)
    return res


def preorder_recursive(root):
    return root and [root.val] + preorder_recursive(root.left) + preorder_recursive(root.right) or []


def inorder_iterative(root):
    stack = []
    res = []
    n = root
    while n or stack:
        if n:
            stack.append(n)
            n = n.left
        else:
            n = stack.pop()
            res.append(n.val)
            n = n.right
    return res


def inorder_recursive(root):
    return root and inorder_recursive(root.left) + [root.val] + inorder_recursive(root.right) or []


def postorder_iterative(root):
    stack = []
    res = []
    n = root
    while n or stack:
        if n:
            stack.append(n)
            stack.append(n)
            n = n.left
        else:
            n = stack.pop()
            if stack and stack[-1] == n:
                n = n.right
            else:
                res.append(n.val)
                n = None
    return res


def postorder_recursive(root):
    return root and postorder_recursive(root.left) + postorder_recursive(root.right) + [root.val] or []


def delete_node(root, vals):
    if root.left:
        if root.left.val in vals:
            root.left = None
        else:
            delete_node(root.left, vals)
    if root.right:
        if root.right.val in vals:
            root.right = None
        else:
            delete_node(root.right, vals)


def level_order_iterative(root):
    if not root:
        return []
    queue = deque((root,))
    res = []
    while queue:
        for _ in range(len(queue)):
            n = queue.popleft()
            res.append(n.val)
            if n.left:
                queue.append(n.left)
            if n.right:
                queue.append(n.right)
    return res


'''
Level order traversal is a kind of BFS. Due to the nature of BFS, there is no
way to traversal a tree in level order recursively, unless modifying the tree.
See more discussions at: 
https://stackoverflow.com/questions/2549541/performing-breadth-first-search-recursively#answer-2550084

This is a tricky method.
'''
def level_order_recursive_modify(root):
    if not root:
        return []
    root.next = None

    def link(prev, child):
        if prev:
            prev.next = child
        child.next = None
        return child

    def helper(node, prev, next_level=True):
        if not node:
            return None, []
        if node.left:
            prev = link(prev, node.left)
        if node.right:
            prev = link(prev, node.right)
        first = node.left if node.left else node.right
        second, res = helper(node.next, prev, next_level=False)
        if next_level:
            _, res2 = helper(first or second, None)
            return _, [node.val] + res + res2
        return first or second, [node.val] + res

    return helper(root, None)[1]


if __name__ == '__main__':
    # build tree
    from collections import deque
    root = TreeNode(1)
    queue = deque((root,))
    for i in range(2, 16):
        if not queue[0].left:
            queue[0].left = TreeNode(i)
            queue.append(queue[0].left)
        else:
            n = queue.popleft()
            n.right = TreeNode(i)
            queue.append(n.right)

    assert preorder_iterative(root) == [1, 2, 4, 8, 9, 5, 10, 11, 3, 6, 12, 13, 7, 14, 15]
    assert preorder_recursive(root) == [1, 2, 4, 8, 9, 5, 10, 11, 3, 6, 12, 13, 7, 14, 15]

    assert inorder_iterative(root) == [8, 4, 9, 2, 10, 5, 11, 1, 12, 6, 13, 3, 14, 7, 15]
    assert inorder_recursive(root) == [8, 4, 9, 2, 10, 5, 11, 1, 12, 6, 13, 3, 14, 7, 15]

    assert postorder_iterative(root) == [8, 9, 4, 10, 11, 5, 2, 12, 13, 6, 14, 15, 7, 3, 1]
    assert postorder_recursive(root) == [8, 9, 4, 10, 11, 5, 2, 12, 13, 6, 14, 15, 7, 3, 1]

    assert level_order_iterative(root) == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
    assert level_order_recursive_modify(root) == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

    delete_node(root, (4, 7, 10, 11, 12))

    assert preorder_iterative(root) == [1, 2, 5, 3, 6, 13]
    assert preorder_recursive(root) == [1, 2, 5, 3, 6, 13]

    assert inorder_iterative(root) == [2, 5, 1, 6, 13, 3]
    assert inorder_recursive(root) == [2, 5, 1, 6, 13, 3]

    assert postorder_iterative(root) == [5, 2, 13, 6, 3, 1]
    assert postorder_recursive(root) == [5, 2, 13, 6, 3, 1]

    assert level_order_iterative(root) == [1, 2, 3, 5, 6, 13]
    assert level_order_recursive_modify(root) == [1, 2, 3, 5, 6, 13]