Invert binary tree

invert_binary_tree.py

class Node:
  def __init__(self, value, left=None, right=None):
    self.value = value
    self.left = left
    self.right = right


def invert_binary_tree(root):
  if root:
    stack = [root]
    
    while stack:
      curr_node = stack.pop()

      curr_node.left, curr_node.right = curr_node.right, curr_node.left

      if curr_node.left:
        stack.append(curr_node.left)

      if curr_node.right:
        stack.append(curr_node.right)

  return root

# Input:
#       1
#      / \
#     2   3
#    /     \
#   4       5

# Vamos representar a árvore acima
# Criando a raíz
root = Node(1)

# Criando os nodos da esquerda
root.left = Node(2)
root.left.left = Node(4)

# Criando os nodos da direita
root.right = Node(3)
root.right.right = Node(5)

print("Before inverting")
print(root.value)
print(root.left.value)
print(root.right.value)
print(root.left.left.value)
print(root.right.right.value)

invert_binary_tree(root)

# Resultado:
#       1
#      / \
#     3   2
#    /     \
#   5       4

print()

print("After inverting")
print(root.value)
print(root.left.value)
print(root.right.value)
print(root.left.left.value)
print(root.right.right.value)