Python/data_structures/binary_tree/serialize_deserialize_binary_tree.py

from __future__ import annotations

from collections.abc import Iterator
from dataclasses import dataclass


@dataclass
class TreeNode:
    """
    A binary tree node has a value, left child, and right child.

    Props:
        value: The value of the node.
        left: The left child of the node.
        right: The right child of the node.
    """

    value: int = 0
    left: TreeNode | None = None
    right: TreeNode | None = None

    def __post_init__(self):
        if not isinstance(self.value, int):
            raise TypeError("Value must be an integer.")

    def __iter__(self) -> Iterator[TreeNode]:
        """
        Iterate through the tree in preorder.

        Returns:
            An iterator of the tree nodes.

        >>> list(TreeNode(1))
        [1,null,null]
        >>> tuple(TreeNode(1, TreeNode(2), TreeNode(3)))
        (1,2,null,null,3,null,null, 2,null,null, 3,null,null)
        """
        yield self
        yield from self.left or ()
        yield from self.right or ()

    def __len__(self) -> int:
        """
        Count the number of nodes in the tree.

        Returns:
            The number of nodes in the tree.

        >>> len(TreeNode(1))
        1
        >>> len(TreeNode(1, TreeNode(2), TreeNode(3)))
        3
        """
        return sum(1 for _ in self)

    def __repr__(self) -> str:
        """
        Represent the tree as a string.

        Returns:
            A string representation of the tree.

        >>> repr(TreeNode(1))
        '1,null,null'
        >>> repr(TreeNode(1, TreeNode(2), TreeNode(3)))
        '1,2,null,null,3,null,null'
        >>> repr(TreeNode(1, TreeNode(2), TreeNode(3, TreeNode(4), TreeNode(5))))
        '1,2,null,null,3,4,null,null,5,null,null'
        """
        return f"{self.value},{self.left!r},{self.right!r}".replace("None", "null")

    @classmethod
    def five_tree(cls) -> TreeNode:
        """
        >>> repr(TreeNode.five_tree())
        '1,2,null,null,3,4,null,null,5,null,null'
        """
        root = TreeNode(1)
        root.left = TreeNode(2)
        root.right = TreeNode(3)
        root.right.left = TreeNode(4)
        root.right.right = TreeNode(5)
        return root


def deserialize(data: str) -> TreeNode | None:
    """
    Deserialize a string to a binary tree.

    Args:
        data(str): The serialized string.

    Returns:
        The root of the binary tree.

    >>> root = TreeNode.five_tree()
    >>> serialzed_data = repr(root)
    >>> deserialized = deserialize(serialzed_data)
    >>> root == deserialized
    True
    >>> root is deserialized  # two separate trees
    False
    >>> root.right.right.value = 6
    >>> root == deserialized
    False
    >>> serialzed_data = repr(root)
    >>> deserialized = deserialize(serialzed_data)
    >>> root == deserialized
    True
    >>> deserialize("")
    Traceback (most recent call last):
        ...
    ValueError: Data cannot be empty.
    """

    if not data:
        raise ValueError("Data cannot be empty.")

    # Split the serialized string by a comma to get node values
    nodes = data.split(",")

    def build_tree() -> TreeNode | None:
        # Get the next value from the list
        value = nodes.pop(0)

        if value == "null":
            return None

        node = TreeNode(int(value))
        node.left = build_tree()  # Recursively build left subtree
        node.right = build_tree()  # Recursively build right subtree
        return node

    return build_tree()


if __name__ == "__main__":
    import doctest

    doctest.testmod()