K^th Largest Element in an Array

Step 1: Sorting

Sort and return nums[n - k]. Simple, O(n log n).

Step 2: A min-heap of size k

Keep the k largest values seen so far in a min-heap. When it grows to k + 1, remove its smallest. At the end, the heap's top is the smallest of the k largest values: the answer. This is O(n log k), good when k is small.

Step 3: Quickselect: sort only what matters

Quick sort's partition step picks a pivot and moves smaller values to its left and larger ones to its right. After that, the pivot is exactly where it would be in the sorted array:

The answer is at index target = n - k of the sorted array. Compare the pivot's position with it:

• equal: the pivot is the answer;
• pivot too far left: the answer is in the right part, so partition only that;
• pivot too far right: partition only the left part.

Unlike quick sort, only one side is ever processed, which makes the average O(n).

nums = [3, 2, 1, 9, 8, 5, 6], k = 2, so target = 7 - 2 = 5.

The pivot is the last element, 6:

store starts at 0. It marks where the next value smaller than the pivot goes:

3, 2 and 1 are each smaller than 6, so each is swapped into position store (here, with itself) and store moves up to 3:

9 and 8 are not smaller, so nothing moves:

5 is smaller, so it swaps with 9 at position 3, and store becomes 4: [3, 2, 1, 5, 8, 9, 6].

Finally the pivot swaps into position store: [3, 2, 1, 5, 6, 9, 8]. 6 is now at index 4, its sorted position.

4 is less than the target 5, so the answer is in the right part [9, 8]. Partitioning that part with pivot 8 puts 8 at index 5, the target. The answer is 8.

Both methods are below: findKthLargestHeap with a PriorityQueue, and findKthLargest with quickselect.

import java.util.PriorityQueue;

public class FindKthLargest {

    // Min-heap of the k largest values seen so far: its top is the k-th largest
    public static int findKthLargestHeap(int[] nums, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>();
        for (int x : nums) {
            heap.offer(x);
            if (heap.size() > k) {
                heap.poll();                  // drop the smallest of the k + 1
            }
        }
        return heap.peek();
    }

    // Quickselect: the k-th largest sits at index n - k once the array is sorted
    public static int findKthLargest(int[] nums, int k) {
        int target = nums.length - k;
        int left = 0;
        int right = nums.length - 1;
        while (true) {
            int p = partition(nums, left, right);
            if (p == target) {
                return nums[p];
            } else if (p < target) {
                left = p + 1;                 // the answer is right of the pivot
            } else {
                right = p - 1;                // the answer is left of the pivot
            }
        }
    }

    // Lomuto partition: smaller values go left of the pivot; returns the pivot's final index
    private static int partition(int[] a, int left, int right) {
        int pivot = a[right];
        int store = left;
        for (int i = left; i < right; i++) {
            if (a[i] < pivot) {
                swap(a, i, store++);
            }
        }
        swap(a, store, right);
        return store;
    }

    private static void swap(int[] a, int i, int j) {
        int tmp = a[i];
        a[i] = a[j];
        a[j] = tmp;
    }

    public static void main(String[] args) {
        System.out.println(findKthLargest(new int[] {3, 2, 1, 5, 6, 4}, 2));              // 5
        System.out.println(findKthLargest(new int[] {3, 2, 3, 1, 2, 4, 5, 5, 6}, 4));     // 4
        System.out.println(findKthLargestHeap(new int[] {3, 2, 1, 9, 8, 5, 6}, 2));       // 8
    }
}

The original implementation and its tests are on GitHub: KthLargestElement.java and TestKthLargestElement.java.

• Sorting: O(n log n) time.
• Heap: O(n log k) time, O(k) space.
• Quickselect: O(n) on average, O(n²) in the worst case (a bad pivot every time, as in an already sorted array). Picking a random pivot makes the worst case very unlikely. O(1) extra space.

The partition step is explained in Quick Sort. The same tools find the most common values in Top K Frequent Elements.