#226. Invert Binary Tree | LeetCode Python Solution

Problem

LeetCode Problem

226. Invert Binary Tree

Python Solution

# 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
class Solution:
    def invertTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
        # Return immediately if root node is missing
        if not root:
            return root

        # Initialize the stack with the root node
        node_stack = [root]

        # Loop until the stack is empty
        while node_stack:
            # Pop the item on the top of the stack
            curr = node_stack.pop()

            # Skip this iteration if this is a leaf node
            if curr.left is None and curr.right is None:
                continue

            # Invert left and right pointer for the current tree node
            tmp = curr.left
            curr.left = curr.right
            curr.right = tmp

            # Push left and right child nodes into the stack so we can invert their right and left pointers
            if curr.left:
                node_stack.append(curr.left)

            if curr.right:
                node_stack.append(curr.right)

        # Return the root which now has an inverted tree
        return root