Update Kth Smallest Element in a BST.java

varunu28 · web-flow · commit b66ddaa489dc · 2022-09-19T10:26:58.000-07:00
diff --git a/Medium/Kth Smallest Element in a BST.java b/Medium/Kth Smallest Element in a BST.java
@@ -15,17 +15,21 @@
  */
 class Solution {
   public int kthSmallest(TreeNode root, int k) {
+    Integer result = null;
     Stack<TreeNode> stack = new Stack<>();
-    while (true) {
-      while (root != null) {
-        stack.push(root);
-        root = root.left;
-      }
-      root = stack.pop();
-      if (--k == 0) {
-        return root.val;
+    while (root != null) {
+      stack.push(root);
+      root = root.left;
+    }
+    while (k > 1 && !stack.isEmpty()) {
+      TreeNode removed = stack.pop();
+      TreeNode rightNode = removed.right;
+      while (rightNode != null) {
+        stack.push(rightNode);
+        rightNode = rightNode.left;
       }
-      root = root.right;
+      k--;
     }
+    return stack.peek().val;
   }
 }
