Find the nTh smallest element in an array

Couple of days ago I came across a problem which asked me to find the nth smallest element in an array. It was marked as a medium level problem, and it quickly got me thinking, why is it so hard??? Simply sort the array and then return the value which is present in the index...

// This is assuming the nthSmallestIndex is within the array.
public int getNthSmallestElement(int[] array, int nthSmallestIndex) {
  Arrays.sort(array);
  return array[nthSmallestIndex];
}

The above code seems quite elementary, so why would this be a medium level problem. Well, my question was answered quite soon. When I ran the code, the online editor complained that my code was slow. To sort an array where the range is completely unknown, the fastest which it can be done is in O(n log n). The question required me to complete the code in O(n) time.

It then occurred to me that I need not sort the entire array. I only need to sort the array up-to the nth smallest index. So if I have the following array [5, 2, 1, 4, 6, 3], and if I need to find the 2nd smallest number, I just need to sort first 3 elements.

Well I could use the lomuto partition to select a pivot element and put it at the correct index and reorder the array where all the elements smaller than the pivot is on left side and greater elements on the right. Let's take a look at how the lomuto partition is implemented.

public int lomutoPartition(int low, int high, int[] array) {
  int pivot = array[high], j = low;
  for (int i = low; i < high; i++) {
    if (array[i] < pivot) {
      int temp = array[i];
      array[i] = array[j];
      array[j] = temp;
      ++j;
    }
  }
  int temp = array[high];
  array[high] = array[j];
  array[j] = temp;
  return j;
}

So if I have my above array, and if I apply the lomuto partition keeping the last element as the pivot, I might get something like this...[2, 1, 3, 6, 4, 5]. Notice how 3 is in the right place and all the elements to the left is lesser than 3 and elements to the right are greater than 3.

Now all I need to check is, if the nth smallest index is greater or smaller than the current pivot and then keep on sorting the array until I reach the required index. Then, it dawned on me, this is QuickSort with a little modification!!! Let's take a look at the code...

  public int findTheNthSmallestElement(int[] array, int nThSmallest, int left, int right) {
    if (left < right) {
      int position = lomutoPartition(left, right, array);
      int count = position - left + 1;
      if (count == nThSmallest) {
        return array[position];
      }
      if (count > nThSmallest) {
        return findTheNthSmallestElement(array, nThSmallest, left, position - 1);
      }
      return findTheNthSmallestElement(array, nThSmallest - count, position + 1, right);
    }
    return -1;
  }

Now this solution on an average case get the nth smallest element in O(n) time. Later when I was doing a google search I found that the above solution is called "QuickSelect"

I hope you this post helped you to understand how to find the nth smallest element in O(n) time (Average case.)

If you would like to see the complete implementation and other suck questions check out my Repo on github.