diff --git a/data_structures/binary_tree/basic_binary_tree.py b/data_structures/binary_tree/basic_binary_tree.py index 9d4c1bdbb57a..aeecbfb927f1 100644 --- a/data_structures/binary_tree/basic_binary_tree.py +++ b/data_structures/binary_tree/basic_binary_tree.py @@ -1,43 +1,70 @@ -from __future__ import annotations +from __future__ import annotations # Allows using the class itself as a type hint within the class -from collections.abc import Iterator -from dataclasses import dataclass +from collections.abc import Iterator # For type hinting iteration +from dataclasses import dataclass # For concise class definitions +# ----------------------------- +# Node class: represents one node in the binary tree +# ----------------------------- @dataclass class Node: - data: int - left: Node | None = None - right: Node | None = None - + data: int # Value stored in the node + left: Node | None = None # Left child node + right: Node | None = None # Right child node + + # ----------------------------- + # In-order traversal generator + # Allows iterating over Node like: `for x in node: ...` + # ----------------------------- def __iter__(self) -> Iterator[int]: if self.left: - yield from self.left - yield self.data + yield from self.left # Yield all values from the left subtree + yield self.data # Yield the current node's data if self.right: - yield from self.right + yield from self.right # Yield all values from the right subtree + # ----------------------------- + # Returns number of nodes in the subtree rooted at this node + # ----------------------------- def __len__(self) -> int: - return sum(1 for _ in self) + return sum(1 for _ in self) # Count all nodes using iteration + # ----------------------------- + # Check if the subtree rooted at this node is a full binary tree + # A full binary tree: every node has 0 or 2 children + # ----------------------------- def is_full(self) -> bool: if not self or (not self.left and not self.right): - return True + return True # Leaf node is full if self.left and self.right: + # Node has two children → recursively check both subtrees return self.left.is_full() and self.right.is_full() - return False + return False # Node has only one child → not full +# ----------------------------- +# BinaryTree class: represents the whole binary tree +# ----------------------------- @dataclass class BinaryTree: - root: Node + root: Node # Root node of the tree + # ----------------------------- + # Iterate over tree using in-order traversal + # ----------------------------- def __iter__(self) -> Iterator[int]: return iter(self.root) + # ----------------------------- + # Return total number of nodes in the tree + # ----------------------------- def __len__(self) -> int: return len(self.root) + # ----------------------------- + # Factory method: small tree (3 nodes) + # ----------------------------- @classmethod def small_tree(cls) -> BinaryTree: """ @@ -48,22 +75,25 @@ def small_tree(cls) -> BinaryTree: >>> list(binary_tree) [1, 2, 3] """ - binary_tree = BinaryTree(Node(2)) - binary_tree.root.left = Node(1) - binary_tree.root.right = Node(3) + binary_tree = BinaryTree(Node(2)) # Root node = 2 + binary_tree.root.left = Node(1) # Left child + binary_tree.root.right = Node(3) # Right child return binary_tree + # ----------------------------- + # Factory method: medium tree (7 nodes) + # ----------------------------- @classmethod def medium_tree(cls) -> BinaryTree: """ - Return a medium binary tree with 3 nodes. + Return a medium binary tree with 7 nodes. >>> binary_tree = BinaryTree.medium_tree() >>> len(binary_tree) 7 >>> list(binary_tree) [1, 2, 3, 4, 5, 6, 7] """ - binary_tree = BinaryTree(Node(4)) + binary_tree = BinaryTree(Node(4)) # Root node = 4 binary_tree.root.left = two = Node(2) two.left = Node(1) two.right = Node(3) @@ -72,10 +102,12 @@ def medium_tree(cls) -> BinaryTree: six.right = Node(7) return binary_tree + # ----------------------------- + # Public method: get depth of tree + # ----------------------------- def depth(self) -> int: """ Returns the depth of the tree - >>> BinaryTree(Node(1)).depth() 1 >>> BinaryTree.small_tree().depth() @@ -85,15 +117,21 @@ def depth(self) -> int: """ return self._depth(self.root) + # ----------------------------- + # Helper recursive method to compute depth + # Depth = 1 + max(depth of left, depth of right) + # ----------------------------- def _depth(self, node: Node | None) -> int: if not node: - return 0 + return 0 # Empty node contributes 0 to depth return 1 + max(self._depth(node.left), self._depth(node.right)) + # ----------------------------- + # Check if the tree is full + # ----------------------------- def is_full(self) -> bool: """ Returns True if the tree is full - >>> BinaryTree(Node(1)).is_full() True >>> BinaryTree.small_tree().is_full() @@ -104,7 +142,10 @@ def is_full(self) -> bool: return self.root.is_full() +# ----------------------------- +# Run doctests if executed as main +# ----------------------------- if __name__ == "__main__": import doctest - doctest.testmod() + doctest.testmod() # Automatically tests all docstring examples