Update basic_binary_tree.py

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 @@
         >>> 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 @@
         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 @@
         """
         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 @@
         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