Python/data_structures/binary_tree/is_sum_tree.py
pre-commit-ci[bot] bc8df6de31
[pre-commit.ci] pre-commit autoupdate ()
* [pre-commit.ci] pre-commit autoupdate

updates:
- [github.com/astral-sh/ruff-pre-commit: v0.2.2 → v0.3.2](https://github.com/astral-sh/ruff-pre-commit/compare/v0.2.2...v0.3.2)
- [github.com/pre-commit/mirrors-mypy: v1.8.0 → v1.9.0](https://github.com/pre-commit/mirrors-mypy/compare/v1.8.0...v1.9.0)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2024-03-13 07:52:41 +01:00

163 lines
4.0 KiB
Python

"""
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 = }.")