Modified inorder traversal to only traverse up to K elements - O(K) solution using stack


  • 0
    I
    public int kthSmallest(TreeNode root, int k) {
       if(k < 1 || root==null) return -1;
       
       Stack<TreeNode> nodeStack = new Stack<>();
       int count = 0;
       TreeNode curr = root;
       while(true){
           while(curr!=null){
               nodeStack.push(curr);
               curr = curr.left;
           }
           curr = nodeStack.pop();
           count++;
           if(count == k) return curr.val;
           
           while(curr.right == null){
               curr = nodeStack.pop();
               count++;
               if(count == k) return curr.val;
           } // end while for curr.right
           curr = curr.right;
       } // end while true;
    }

  • 0
    G

    Isn't this O(h + k) (where h is the height)?


  • 0
    I

    No, How did the height of the tree come in the picture?
    If you look at the code, all it is doing is simulating in-order traversal with a stack and keeps counting how many nodes have been traversed in-order. When count reaches K, i break and return that value.
    Makes sense?


  • 0
    G

    Doesn't your while(curr!=null) loop take O(h) steps? As I see your code, it first reaches the leftmost node in O(n) steps, then it walks the tree finding the desired node in O(k) steps.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.