Finding median or Kth element in Binary Search Tree (BST)

Finding kth element in BST is quite simple. We know that all the elements in left subtree is smaller than all the elements in right subtree. If we count the number of elements in left and right subtree then we know in which subtree the required element exists and then we move to this subtree and repeat the same steps again. For example starting from root we are looking for 11th element. Lets say left subtree has 7 elements and right subtree has 6 elements.Then as the number 11 is greater than 7 we know that the 11th median exists in the right subtree. Now in the right subtree we look for (11 - 7 -1 = 3)rd element (Why?? because we have discarded left subtree). We proceed similarly in the new subtree until we get the kth median as the root.

The example described above can very well be described in recursive form. The code below is a complete program in recursive form.

  
 //find kth element in a BST recursively
Node * find_kth_element(Node *tree, int &k) {

 if(tree == 0)
  return 0;

 int count = 0;

 //count the nodes in left subtree
 count_nodes(tree->left_, count);

 //check if the root is median
 if( k == count + 1)
  return tree;

 //check if median falls on left subtree
 else if (count >= k) {

  return find_kth_element(tree->left_,k);
 }
 //the median falls on right subtree
 else {
  k = k - (count + 1);
  return find_kth_element(tree->right_, k);
 }
}

Iterative Solution

We know that recursive solutions are not that efficient and it's always good to avoid it if it's possible. The problem however is equally simple in iterative form as well. The code below translates recursive form into iterative form.

  
 

//find kth element iteratively
Node * find_kth_element_iterative(Node *tree, int &k) {
 if(tree == 0)
  return 0;

 Node *node = tree;
 int count = 0;
 int pos = k;


 while(node != 0) {
  count = 0;

  //count nodes on the left subtree
  count_nodes(node->left_, count);

  //check if root is the median
  if( pos == count + 1)
   return node;

  //check if median falls on left subtree
  else if (count >= pos)
   node = node->left_;

  //median falls on right subtree
  else {
   pos -= (count + 1);
   node = node->right_;
  }

 }
 return 0;
}

Complexity

What's the complexity of above algorithms? While finding median You are moving along a path in the tree, so in the worst case it can be height of the tree (h) and also you are counting the number of nodes in left subtree at each step. Which requires O(n) time in the worst case. Thus the overall algorithm is O(nh).

Is there a way to improve it? Yes, of course ( at least one technique I know of). If you need to keep on computing the kth median on a single tree many times, then you can construct a tree with a count of elements in its subtree. If you do this the counting of number of nodes in left subtree takes O(1) time. Effectively the algorithm will reduce to O(h) algorithm. In case tree is balanced this is just O(log n). Not bad :).