Invert Binary Tree Python Solution with Tests

Tested Python solution for LeetCode 226 with 15 pytest cases. Generate a practice environment with lcpy.

LeetCode 226, Easy. Topics: Tree, Depth-First Search, Breadth-First Search, Binary Tree. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 15 parametrized pytest cases, and a playground notebook:

lcpy gen -n 226   # by problem number
lcpy gen -s invert_binary_tree   # by problem name

​

Problem

Given the root of a binary tree, invert the tree, and return its root.

​

Examples

Input: root = [4,2,7,1,3,6,9]
Output: [4,7,2,9,6,3,1]

Input: root = [2,1,3]
Output: [2,3,1]

Input: root = []
Output: []

​

Constraints

• The number of nodes in the tree is in the range [0, 100]
• -100 <= Node.val <= 100

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

from collections import deque

from leetcode_py import TreeNode

# Note: "Fringe" is the general CS term for the data structure holding nodes to be explored.
# Stack (LIFO) → DFS, Queue (FIFO) → BFS, Priority Queue → A*/Best-first search


class Solution:
    # DFS recursive
    # Time: O(n)
    # Space: O(h) where h is height of tree
    def invert_tree(self, root: TreeNode[int] | None) -> TreeNode[int] | None:
        if not root:
            return None

        root.left, root.right = self.invert_tree(root.right), self.invert_tree(root.left)
        return root


class SolutionDFS:
    # DFS iterative
    # Time: O(n)
    # Space: O(h) where h is height of tree
    def invert_tree(self, root: TreeNode[int] | None) -> TreeNode[int] | None:
        if not root:
            return None

        stack: list[TreeNode[int] | None] = [root]
        while stack:
            node = stack.pop()
            if node is None:
                continue
            node.left, node.right = node.right, node.left

            stack.append(node.left)
            stack.append(node.right)

        return root


class SolutionBFS:
    # Time: O(n)
    # Space: O(w) where w is maximum width of tree
    def invert_tree(self, root: TreeNode[int] | None) -> TreeNode[int] | None:
        if not root:
            return None

        queue: deque[TreeNode[int] | None] = deque([root])
        while queue:
            node = queue.popleft()
            if node is None:
                continue
            node.left, node.right = node.right, node.left

            queue.append(node.left)
            queue.append(node.right)

        return root

​

Complexity

Time	Space
O(n)	O(h) where h is height of tree

​

Tags

Grind 75, Grind, Blind 75, NeetCode 150, NeetCode 250, NeetCode All.