BSTIterator

Author: volodimyr

173.binary_search_tree_iterator/iterator.py

@@ -0,0 +1,27 @@
+from typing import Optional
+
+class BSTIterator:
+    def __init__(self, root: Optional[TreeNode]):
+        self.stack = []
+        self.root = root
+        cur = root
+        while cur:
+            self.stack.append(cur)
+            cur = cur.left
+
+    def next(self) -> int:
+        pop = self.stack.pop()
+        cur = pop.right
+        while cur:
+            self.stack.append(cur)
+            cur = cur.left
+        return pop.val
+
+    def hasNext(self) -> bool:
+        return len(self.stack) > 0
+    
+class TreeNode:
+    def __init__(self, val:0, left: None, right: None):
+        self.val = val
+        self.left = left
+        self.right = right

173.binary_search_tree_iterator/test_iterator.py

@@ -0,0 +1,56 @@
+import unittest
+from iterator import TreeNode
+from iterator import BSTIterator
+
+class TestBSTIterator(unittest.TestCase):
+    def build_tree(self):
+        # Build a BST:
+        #        7
+        #      /   \
+        #     3     15
+        #          /  \
+        #         9    20
+        n9 = TreeNode(9, None, None)
+        n20 = TreeNode(20, None, None)
+        n15 = TreeNode(15, n9, n20)
+        n3 = TreeNode(3, None, None)
+        root = TreeNode(7, n3, n15)
+        return root
+
+    def test_inorder_traversal(self):
+        root = self.build_tree()
+        it = BSTIterator(root)
+        result = []
+        while it.hasNext():
+            result.append(it.next())
+        self.assertEqual(result, [3, 7, 9, 15, 20])
+
+    def test_single_node(self):
+        root = TreeNode(42, None, None)
+        it = BSTIterator(root)
+        self.assertTrue(it.hasNext())
+        self.assertEqual(it.next(), 42)
+        self.assertFalse(it.hasNext())
+
+    def test_empty_tree(self):
+        it = BSTIterator(None)
+        self.assertFalse(it.hasNext())
+
+    def test_alternating_next_and_hasNext(self):
+        root = self.build_tree()
+        it = BSTIterator(root)
+        self.assertTrue(it.hasNext())
+        self.assertEqual(it.next(), 3)
+        self.assertTrue(it.hasNext())
+        self.assertEqual(it.next(), 7)
+        self.assertTrue(it.hasNext())
+        self.assertEqual(it.next(), 9)
+        self.assertTrue(it.hasNext())
+        self.assertEqual(it.next(), 15)
+        self.assertTrue(it.hasNext())
+        self.assertEqual(it.next(), 20)
+        self.assertFalse(it.hasNext())
+
+
+if __name__ == '__main__':
+    unittest.main()