Inorder, Preorder and postorder iterations

Author: volodimyr

173.binary_search_tree_iterator/iterator.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import List, Optional
 
 class BSTIterator:
     def __init__(self, root: Optional[TreeNode]):
@@ -21,7 +21,50 @@ def hasNext(self) -> bool:
         return len(self.stack) > 0
 
 class TreeNode:
-    def __init__(self, val:0, left: None, right: None):
+    def __init__(self, val=0, left=None, right=None):
         self.val = val
         self.left = left
-        self.right = right
+        self.right = right
+
+    def inorder(self) -> List[int]:
+        cur = self
+        stack = []
+        result = []
+        while cur or stack:
+            if cur:
+                stack.append(cur)
+                cur = cur.left
+            else:
+                pop = stack.pop()
+                result.append(pop.val)
+                cur = pop.right
+        return result
+
+    def preorder(self) -> List[int]:
+        cur = self
+        stack = []
+        result = []
+        while cur or stack:
+            if cur:
+                result.append(cur.val)
+                stack.append(cur)
+                cur = cur.left
+            else:
+                pop = stack.pop()
+                cur = pop.right
+        return result
+
+    def postorder(self) -> List[int]:
+        stack = []
+        stack.append((self, False))
+        result = []
+        while stack:
+            cur = stack.pop()
+            if cur[0]:
+                if cur[1]:
+                    result.append(cur[0].val)
+                else:
+                    stack.append((cur[0], True))
+                    stack.append((cur[0].right, False))
+                    stack.append((cur[0].left, False))
+        return result

173.binary_search_tree_iterator/test_iterator.py

@@ -2,6 +2,118 @@
 from iterator import TreeNode
 from iterator import BSTIterator
 
+
+class TestTreeNodePostorder(unittest.TestCase):
+    def test_single_node(self):
+        root = TreeNode(1)
+        self.assertEqual(root.postorder(), [1])
+
+    def test_left_skewed(self):
+        #   3
+        #  /
+        # 2
+        #/
+        #1
+        root = TreeNode(3, TreeNode(2, TreeNode(1)))
+        self.assertEqual(root.postorder(), [1, 2, 3])
+
+    def test_right_skewed(self):
+        # 1
+        #  \
+        #   2
+        #    \
+        #     3
+        root = TreeNode(1, None, TreeNode(2, None, TreeNode(3)))
+        self.assertEqual(root.postorder(), [3, 2, 1])
+
+    def test_balanced_tree(self):
+        #   2
+        #  / \
+        # 1   3
+        root = TreeNode(2, TreeNode(1), TreeNode(3))
+        self.assertEqual(root.postorder(), [1, 3, 2])
+
+    def test_complex_tree(self):
+        #        4
+        #      /   \
+        #     2     6
+        #    / \   /
+        #   1   3 5
+        root = TreeNode(4,
+                        TreeNode(2, TreeNode(1), TreeNode(3)),
+                        TreeNode(6, TreeNode(5)))
+        self.assertEqual(root.postorder(), [1, 3, 2, 5, 6, 4])
+
+
+class TestTreeNodePreorder(unittest.TestCase):
+    def test_single_node(self):
+        root = TreeNode(1)
+        self.assertEqual(root.preorder(), [1])
+
+    def test_left_skewed(self):
+        #   3
+        #  /
+        # 2
+        #/
+        #1
+        root = TreeNode(3, TreeNode(2, TreeNode(1)))
+        self.assertEqual(root.preorder(), [3, 2, 1])
+
+    def test_right_skewed(self):
+        # 1
+        #  \
+        #   2
+        #    \
+        #     3
+        root = TreeNode(1, None, TreeNode(2, None, TreeNode(3)))
+        self.assertEqual(root.preorder(), [1, 2, 3])
+
+    def test_balanced_tree(self):
+        #   2
+        #  / \
+        # 1   3
+        root = TreeNode(2, TreeNode(1), TreeNode(3))
+        self.assertEqual(root.preorder(), [2, 1, 3])
+
+    def test_complex_tree(self):
+        #        4
+        #      /   \
+        #     2     6
+        #    / \   /
+        #   1   3 5
+        root = TreeNode(4,
+                        TreeNode(2, TreeNode(1), TreeNode(3)),
+                        TreeNode(6, TreeNode(5)))
+        self.assertEqual(root.preorder(), [4, 2, 1, 3, 6, 5])
+
+class TestTreeNode(unittest.TestCase):
+    def test_single_node(self):
+        root = TreeNode(1, None, None)
+        self.assertEqual(root.inorder(), [1])
+
+    def test_left_skewed(self):
+        root = TreeNode(3, TreeNode(2, TreeNode(1)))
+        self.assertEqual(root.inorder(), [1, 2, 3])
+
+    def test_right_skewed(self):
+        root = TreeNode(1, None, TreeNode(2, None, TreeNode(3)))
+        self.assertEqual(root.inorder(), [1, 2, 3])
+
+    def test_balanced_tree(self):
+        root = TreeNode(2, TreeNode(1), TreeNode(3))
+        self.assertEqual(root.inorder(), [1, 2, 3])
+
+    def test_complex_tree(self):
+        #        4
+        #      /   \
+        #     2     6
+        #    / \   /
+        #   1   3 5
+        root = TreeNode(4,
+                        TreeNode(2, TreeNode(1), TreeNode(3)),
+                        TreeNode(6, TreeNode(5)))
+        self.assertEqual(root.inorder(), [1, 2, 3, 4, 5, 6])
+
 class TestBSTIterator(unittest.TestCase):
     def build_tree(self):
         # Build a BST: