226. Invert Binary Tree

Invert a binary tree.

Example:

Input:

     4
   /   \
  2     7
 / \   / \
1   3 6   9

Output:

     4
   /   \
  7     2
 / \   / \
9   6 3   1

Source

In [1]:

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

Code

In [2]:

def invert_tree(root):
    """Recursive approach"""
    if root and (root.left or root.right):
        root.left, root.right = root.right, root.left
        invert_tree(root.left)
        invert_tree(root.right)
    return root

Follow up:

Solve it both recursively and iteratively.

Code

In [8]:

def invert_tree(root):
    """Iterative approach"""
    if not root:
        return None

    stack = [root]
    while stack:
        node = stack.pop()
        if node:
            node.left, node.right = node.right, node.left
            stack.append(node.left)
            stack.append(node.right)
    return root