Inverting a binary tree in Python (Leetcode)

bintree_leetcode.py

# See https://leetcode.com/problems/invert-binary-tree/
# for brief and inputs

from collections import deque
from typing import Union, cast


# Note that we use the Union type rather than the pipe syntax
# as the version of Python in this project's pyenv
# doesn't support recursive types using the latter.
class Node:
    left: Union["Node", None] = None
    right: Union["Node", None] = None

    def __init__(self, value: int):
        self.value = value


def build_tree(tree_rep: deque[int]) -> Node | None:
    root = Node(tree_rep.popleft())
    next_parents = deque[Node]([root])

    while len(tree_rep):
        left_val = tree_rep.popleft()
        right_val = tree_rep.popleft()

        parent = next_parents.popleft()
        parent.left = Node(left_val) if left_val else None
        parent.right = Node(right_val) if right_val else None

        if parent.left:
            next_parents.append(parent.left)

        if parent.right:
            next_parents.append(parent.right)

    return root


def serialise_tree(tree: Node) -> list[int]:
    result: list[int] = []

    result.append(tree.value)
    next_parents = deque([tree])

    while len(next_parents):
        root = next_parents.popleft()

        if root.left:
            next_parents.append(root.left)
            result.append(root.left.value)

        if root.right:
            next_parents.append(root.right)
            result.append(root.right.value)

    return result


def invert(tree_rep: list[int]) -> list[int]:
    tree = build_tree(deque(tree_rep))
    next_parents: list[Node] = [cast(Node, tree)]

    while len(next_parents):
        parent = next_parents.pop()
        left = parent.left
        right = parent.right

        if left:
            next_parents.append(left)
            parent.right = left

        if right:
            next_parents.append(right)
            parent.left = right

    return serialise_tree(cast(Node, tree))


input = [4, 2, 7, 1, 3, 6, 9]
expected = [4, 7, 2, 9, 6, 3, 1]
actual = invert(input)

assert actual == expected, f"actual {actual} != expected {expected}"
print("Tree has been inverted!")