Python/data_structures/binary_tree/is_sum_tree.py

"""
Is a binary tree a sum tree where the value of every non-leaf node is equal to the sum
of the values of its left and right subtrees?
https://www.geeksforgeeks.org/check-if-a-given-binary-tree-is-sumtree
"""

from __future__ import annotations

from collections.abc import Iterator
from dataclasses import dataclass


@dataclass
class Node:
    data: int
    left: Node | None = None
    right: Node | None = None

    def __iter__(self) -> Iterator[int]:
        """
        >>> root = Node(2)
        >>> list(root)
        [2]
        >>> root.left = Node(1)
        >>> tuple(root)
        (1, 2)
        """
        if self.left:
            yield from self.left
        yield self.data
        if self.right:
            yield from self.right

    def __len__(self) -> int:
        """
        >>> root = Node(2)
        >>> len(root)
        1
        >>> root.left = Node(1)
        >>> len(root)
        2
        """
        return sum(1 for _ in self)

    @property
    def is_sum_node(self) -> bool:
        """
        >>> root = Node(3)
        >>> root.is_sum_node
        True
        >>> root.left = Node(1)
        >>> root.is_sum_node
        False
        >>> root.right = Node(2)
        >>> root.is_sum_node
        True
        """
        if not self.left and not self.right:
            return True  # leaf nodes are considered sum nodes
        left_sum = sum(self.left) if self.left else 0
        right_sum = sum(self.right) if self.right else 0
        return all(
            (
                self.data == left_sum + right_sum,
                self.left.is_sum_node if self.left else True,
                self.right.is_sum_node if self.right else True,
            )
        )


@dataclass
class BinaryTree:
    root: Node

    def __iter__(self) -> Iterator[int]:
        """
        >>> list(BinaryTree.build_a_tree())
        [1, 2, 7, 11, 15, 29, 35, 40]
        """
        return iter(self.root)

    def __len__(self) -> int:
        """
        >>> len(BinaryTree.build_a_tree())
        8
        """
        return len(self.root)

    def __str__(self) -> str:
        """
        Returns a string representation of the inorder traversal of the binary tree.

        >>> str(list(BinaryTree.build_a_tree()))
        '[1, 2, 7, 11, 15, 29, 35, 40]'
        """
        return str(list(self))

    @property
    def is_sum_tree(self) -> bool:
        """
        >>> BinaryTree.build_a_tree().is_sum_tree
        False
        >>> BinaryTree.build_a_sum_tree().is_sum_tree
        True
        """
        return self.root.is_sum_node

    @classmethod
    def build_a_tree(cls) -> BinaryTree:
        r"""
        Create a binary tree with the specified structure:
              11
           /     \
          2       29
         / \     /  \
        1   7  15    40
                       \
                        35
        >>> list(BinaryTree.build_a_tree())
        [1, 2, 7, 11, 15, 29, 35, 40]
        """
        tree = BinaryTree(Node(11))
        root = tree.root
        root.left = Node(2)
        root.right = Node(29)
        root.left.left = Node(1)
        root.left.right = Node(7)
        root.right.left = Node(15)
        root.right.right = Node(40)
        root.right.right.left = Node(35)
        return tree

    @classmethod
    def build_a_sum_tree(cls) -> BinaryTree:
        r"""
        Create a binary tree with the specified structure:
             26
            /  \
          10    3
         /  \    \
        4    6    3
        >>> list(BinaryTree.build_a_sum_tree())
        [4, 10, 6, 26, 3, 3]
        """
        tree = BinaryTree(Node(26))
        root = tree.root
        root.left = Node(10)
        root.right = Node(3)
        root.left.left = Node(4)
        root.left.right = Node(6)
        root.right.right = Node(3)
        return tree


if __name__ == "__main__":
    from doctest import testmod

    testmod()
    tree = BinaryTree.build_a_tree()
    print(f"{tree} has {len(tree)} nodes and {tree.is_sum_tree = }.")
    tree = BinaryTree.build_a_sum_tree()
    print(f"{tree} has {len(tree)} nodes and {tree.is_sum_tree = }.")