Solution: Keeping Only the k Largest

The last attempt kept all n elements in the heap, but we only need the top k. So, as we go through the array, let’s keep only the k largest elements seen so far. Once we have seen everything, the smallest of those k is the kth largest overall.

Keeping only k means that the element we remove is the smallest one, not the largest. So we switch from a max-heap to a min-heap: each time we add an element, if the heap then holds more than k elements, its smallest cannot be among the k largest, so we remove it. A min-heap makes that cheap, since it keeps its smallest element at the root.

public static int findKthLargest(int[] numbers, int k) {
  PriorityQueue<Integer> heap = new ArrayHeap<>();
  for (int number : numbers) {
    heap.add(number);
    if (heap.size() > k) {
      heap.remove();  // drop the smallest; it is not among the k largest
    }
  }
  return heap.peek();  // the smallest of the k largest
}

Let’s trace it on numbers = [3, 1, 5, 2, 4] with k = 2. Once it fills, the heap holds k = 2 elements, plus one more for a moment after each add. As before, the diagrams show the contents smallest-first, not in the heap’s array order.

The first two elements just fill the heap. We add 3, then 1; the size is still within k = 2, so nothing is removed.

┌───┐
│ 3 │   size = 1
└───┘
  ▲
 peek → 3   (min)

┌───┬───┐
│ 1 │ 3 │   size = 2
└───┴───┘
  ▲
 peek → 1   (min)

We add 5. The heap now holds three elements, one too many. peek points at 1, which cannot be among the two largest, so we remove it.

┌───┬───┬───┐
│ 1 │ 3 │ 5 │   size = 3
└───┴───┴───┘
  ▲
 peek → 1   (min)

┌───┬───┐
│ 3 │ 5 │   size = 2
└───┴───┘
  ▲
 peek → 3   (min)

We add 2. The size is three again, and peek points at 2, so we remove it.

┌───┬───┬───┐
│ 2 │ 3 │ 5 │   size = 3
└───┴───┴───┘
  ▲
 peek → 2   (min)

┌───┬───┐
│ 3 │ 5 │   size = 2
└───┴───┘
  ▲
 peek → 3   (min)

We add the last element, 4. The size is three once more, so we remove the smallest, 3.

┌───┬───┬───┐
│ 3 │ 4 │ 5 │   size = 3
└───┴───┴───┘
  ▲
 peek → 3   (min)

┌───┬───┐
│ 4 │ 5 │   size = 2
└───┴───┘
  ▲
 peek → 4   (min)

The heap ends holding {4, 5}, the two largest elements, and peek returns 4, the second largest.

What it costs

The loop runs n times, and each pass does an add, a size, and sometimes a remove. The heap never holds more than k + 1 elements, so its height is about $lo g k$ , and each add is $O (lo g k)$ amortized and each remove is $O (lo g k)$ . The whole solution is $O (n lo g k)$ . It also needs only $O (k)$ space, not $O (n)$ , and it leaves the caller’s array unchanged.

At $O (n lo g k)$ this can beat sorting’s $O (n lo g n)$ when k is much smaller than n. If k grows in proportion to n, the two have the same asymptotic order.

For the kth smallest element, flip the order: keep a size-k max-heap, built by passing a reversed comparator, and its peek is the kth smallest.