Median and Order statistics

Finding kth median in linear time

The idea is to partition the array we do in quick sort.  Assuming a 0 based index the item that you are looking for is positioned at k-1  if the array is ordered. Here is how we can find that element. First partition the array using a pivot (the pivot partitions the array into two parts; the elements on the left are smaller than the pivot element and elements on the right are greater than the pivot element).  Lets say the pivot index is i. if i equals k-1 then it is your median. Otherwise check if k-1 is less than  i. If it is  then repeat the same process for elements from 0 to i - 1 else repeat the process for elements from i + 1 to the end of partition that you are working on. As we chose the pivot randomly this algorithm would take expected linear time.

Here is the code that implements the above algorithm.

int find_kth_median(int x[],int start, int end, int k) {

 int pivot =  start + ((float)rand() - 1)/INT_MAX * (end - start +1);
 int a = x[pivot];

 int up = end, down = start;
 while(down < up) {

  while(x[down] <= a && down <= end)
   down++;
  while(x[up] > a)
   up--;
  if(down < up) {
   std::swap(x[up], x[down]);
  }
 }

 x[start] = x[up];
 x[up] = a;

 if (up == k-1)
  return x[up];
 else if (up < k-1) 
  return find_kth_median(x,up+1, end, k);
 else
  return find_kth_median(x, start, up -1, k);
}

Here are some of the questions that are similar to the above question:

Given an algorithm that finds median in linear time. How would you find kth element ?

Use the technique above. The idea is simple: find the median of the whole array ( the array will be partitioned around the media). Say the median index is i. Then if k is greater than i repeat the process on the right side of the array else repeat on the left side. Continue until you find the median.

Given a set of elements divide the set into k equal parts

The idea is to find the median of the whole set and then find median of the set at the left and right of the median and then continue so on. This requires O (n log k).

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 :).