efficiently find kth smallest element in binary search tree?

Question

I was faced with below interview question today and I came up with recursive and iterative solution as mentioned below but somehow interviewer was not happy. I am not sure why.

Given a binary search tree, find the kth smallest element in it.

Is there any better way to do this problem either in recursive or iterative way?

  /*****************************************************
   *
   * Kth Smallest Recursive
   *
   ******************************************************/
  public int kthSmallestRecursive(TreeNode root, int k) {
    int count = countNodes(root.left);
    if (k <= count) {
      return kthSmallestRecursive(root.left, k);
    } else if (k > count + 1) {
      return kthSmallestRecursive(root.right, k - 1 - count);
    }

    return root.data;
  }

  public int countNodes(TreeNode n) {
    if (n == null)
      return 0;

    return 1 + countNodes(n.left) + countNodes(n.right);
  }

  /*****************************************************
   *
   * Kth Smallest Iterative
   *
   ******************************************************/

  public int kthSmallestIterative(TreeNode root, int k) {
    Stack<TreeNode> st = new Stack<>();

    while (root != null) {
      st.push(root);
      root = root.left;
    }

    while (k != 0) {
      TreeNode n = st.pop();
      k--;
      if (k == 0)
        return n.data;
      TreeNode right = n.right;
      while (right != null) {
        st.push(right);
        right = right.left;
      }
    }

    return -1;
  }

I mentioned complexity on both the above solution as O(depth of node) which is O(log n).

My iterative version takes extra space. Is there any way to do it without any extra space?

Answer 1

First, I doubt if it's going to take O(logn). It's at least O(n) for vanilla binary search tree.

For vanilla bst you can do iterative or recursive, both are simulation the same process of ignorer traversal with worst time complexity O(n) and space complexity O(logn) (note that even the recursive solution can take O(logn) space from stack).

However you can speed things up or use less space if you could change the data structure a bit:

1. If you are using a binary search tree with each node recording the number of total nodes under it, you can do it as if it is a sorted array and get to the k-th element with O(logn) time and O(1) space.

2. If you are using a binary search tree with each node containing the parent pointer and a status indicator you can traverse the tree in the same fashion as the vanilla solution. On your way you mark whether the node is a. not traversed 2. has its left node traversed 3. has both nodes traversed. then you can do it in O(n) time and O(1) space.

Answer 2

 public int printInorder(Node node, int k) 
    { 
        if (node == null || k <= 0) //Stop traversing once you found the k-th smallest element
            return k; 

        /* first recur on left child */
        k = printInorder(node.left, k); 

        k--;
        if(k == 0) {  
            System.out.print(node.key);
        }

        /* now recur on right child */
        return printInorder(node.right, k);
    } 

In countNodes() method of your recursive approach, you're traversing the Tree at each and every node for counting. countNodes() operation is O(n) and you're performing the operation on log(n) number of nodes. So the order of complexity of your recursive solution is O(n^2) whereas the approach I mentioned above solves in O(n).