Kth Smallest Element in a BST - Iterative: Morris Traversal

Author: dksifoua

src/main/java/io/dksifoua/leetcode/kthsmallestelementinabst/Solution.java

@@ -22,4 +22,31 @@ private void dfs(TreeNode node, int[] smallest) {
 
         if (node.getRight() != null) dfs(node.getRight(), smallest);
     }
+
+    public int kthSmallestIterative(TreeNode root, int k) {
+        TreeNode current = root;
+        while (current != null) {
+            if (current.getLeft() == null) {
+                k -= 1;
+                if (k == 0) return current.getValue();
+                current = current.getRight();
+            } else {
+                TreeNode previous = current.getLeft();
+                while (previous.getRight() != null && previous.getRight() != current) {
+                    previous = previous.getRight();
+                }
+
+                if (previous.getRight() == null) {
+                    previous.setRight(current);
+                    current = current.getLeft();
+                } else {
+                    k -= 1;
+                    if (k == 0) return current.getValue();
+                    current = current.getRight();
+                }
+            }
+        }
+
+        throw new IllegalArgumentException("Not solvable!!");
+    }
 }

src/test/java/io/dksifoua/leetcode/kthsmallestelementinabst/SolutionTest.java

@@ -3,7 +3,7 @@
 import io.dksifoua.leetcode.utils.TreeNode;
 import org.junit.jupiter.api.Test;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.*;
 
 public class SolutionTest {
 
@@ -14,8 +14,26 @@ void test1() {
         assertEquals(1, solution.kthSmallest(TreeNode.build(new Integer[] { 3, 1, 4, null, 2 }), 1));
     }
 
+    @Test
+    void test1Iterative() {
+        assertEquals(1, solution.kthSmallestIterative(TreeNode.build(new Integer[] { 3, 1, 4, null, 2 }), 1));
+    }
+
     @Test
     void test2() {
         assertEquals(3, solution.kthSmallest(TreeNode.build(new Integer[] { 5, 3, 6, 2, 4, null, null, 1 }), 3));
     }
+
+    @Test
+    void test2Iterative() {
+        assertEquals(3, solution.kthSmallestIterative(TreeNode.build(new Integer[] { 5, 3, 6, 2, 4, null, null, 1 }), 3));
+    }
+
+    @Test
+    void test3Iterative() {
+        Throwable exception = assertThrows(IllegalArgumentException.class, () -> {
+            solution.kthSmallestIterative(TreeNode.build(new Integer[] { 5, 3, 6, 2, 4, null, null, 1 }), 7);
+        });
+        assertEquals("Not solvable!!", exception.getMessage());
+    }
 }