Kth maximum element in unsorted collection

Finding the Kth maximum element in an unsorted array doesn’t seem to be a big deal. Because, we know that finding the maximum and if we keep on eliminating the largest element from the array. And this we will take us to the Kth maximum element in an array. But this solution is of order O(n²). Can it be done in linear way?

Yes. Well, almost.

Have a look at the following method written in Java.

public static int findKthMax(List input, int k) {
	if (input != null && input.size() > 0) {
		int element = input.get(0);
		List largerThanElement = new ArrayList();
		List smallerThanElement = new ArrayList();

		for (int i = 1; i < input.size(); i++) {
			if (element < input.get(i)) {
				largerThanElement.add(input.get(i));
			} else {
				smallerThanElement.add(input.get(i));
			}
		}

		if (largerThanElement.size() == k - 1) {
			return element;
		} else if (largerThanElement.size() < k - 1) {
			return findKthMax(largerThanElement, k - largerThanElement.size() - 1);
		} else if (largerThanElement.size() >= k) {		 
			return findKthMax(largerThanElement, k);
		}
		return element; // To satisfy the crazy compiler
	} else {
		throw new IllegalArgumentException();
	}
}

Generalizing into an algorithm

Calling FindKthMax(int[] input, int k)

1. Pick randomly a pivot a from input, lets call the selected pivot element – a.

2. Partition the n numbers into two sets:

* S – all the numbers smaller than a

* L – all the numbers larger than a

3. If |L| = k – 1 then a is the required K-median. Return a.

4. If |L| < k – 1 then the K-median lies somewhere in S. Call recursively to FindKthMax( S, k- |L| – 1 )

5. Else, call recursively  FindKthMax( L, k ).

Extending the solution to solve other problems:

The algorithm can be extended to find the median as well.
Q. Find the median of an array A.
A. If length of A is odd, the median would be FindKthMedian(a.length/2 + 1, input).
If the length of A is even, the median [FindKthMedian(a.length/2, input) + FindKthMedian(a.length/2 + 1, input)] / 2

And the problems are beautifully solved in linear order on n.