☕ DSA Java Articles

Topics › Heaps › Medium Problems

K-th Largest element in an array

Medium TUF ID 578 A2ZSDE

LeetCode ↗ Java source ↗ Edit this article ↗

Problem

Given an integer array nums and an integer k, return the k-th largest element. This is the k-th largest in sorted order, not the k-th largest distinct value: duplicates each occupy their own position. Do not modify the caller’s array.

Constraints: 1 ≤ k ≤ n ≤ 10^5, values fit in an int (negatives allowed). The LeetCode version asks whether it can be done without sorting.

Example 1

nums = [3, 2, 1, 5, 6, 4], k = 2
answer = 5               // sorted descending: 6, 5, 4, ...; the 2nd is 5

Example 2

nums = [3, 2, 3, 1, 2, 4, 5, 5, 6], k = 4
answer = 4               // descending: 6, 5, 5, 4, ...; both 5s count

Approach 1: Brute force, sort

Intuition. After sorting ascending, the largest element is at index n - 1, the second largest at n - 2, and the k-th largest at n - k. Sorting a copy keeps the input intact.

Algorithm.

  1. Copy nums and sort the copy ascending.
  2. Return copy[n - k].
static int bruteForce(int[] nums, int k) {
    int[] a = nums.clone();
    Arrays.sort(a);
    return a[a.length - k];
}

Complexity. Time O(n log n) for the sort. Space O(n) for the copy.

Sorting orders all n elements, but we only need to know which value lands in one position. Both remaining approaches do less ordering.

Approach 2: Better, min-heap of size k

Intuition. Scan the array while keeping the k largest values seen so far. Which of those k is the most useful to see instantly? The smallest of them: it is the gatekeeper. A new value only deserves a place in the top k if it beats the gatekeeper, and when it does, the gatekeeper is the one that gets evicted. A min-heap keeps the smallest of its elements at the root, so a min-heap of size k is exactly this structure. After the scan the heap holds the k largest values of the whole array, and its root, the smallest of them, is the k-th largest.

Algorithm.

  1. Create an empty min-heap.
  2. For each x in nums: add x; if the heap now has k + 1 elements, remove the root.
  3. Return the root.
static int better(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 k + 1 candidates
    }
    return heap.peek();
}

Dry run on Example 1, k = 2:

x heap after offer size > k? heap after
3 {3} no {3}
2 {2, 3} no {2, 3}
1 {1, 2, 3} yes, drop 1 {2, 3}
5 {2, 3, 5} yes, drop 2 {3, 5}
6 {3, 5, 6} yes, drop 3 {5, 6}
4 {4, 5, 6} yes, drop 4 {5, 6}

The root is 5.

Complexity. Time O(n log k): each of the n elements causes at most one insertion and one removal on a heap of size at most k + 1. Space O(k). This is the approach to reach for when k is small, when memory is tight, or when the numbers arrive as a stream that cannot be stored.

Approach 3: Optimal, quickselect

Intuition. Quicksort’s partition step puts a pivot in its final sorted position, with smaller values on its left and larger values on its right. If we want the value at sorted index target = n - k, after one partition we know which side it lives on, so we recurse into one side only and discard the other. With a random pivot each partition throws away a constant fraction of the range on average, so the work is about n + n/2 + n/4 + ... = 2n, linear on average.

Two details make it robust:

Algorithm.

  1. Copy the array. Set target = n - k, lo = 0, hi = n - 1.
  2. Repeat: pick a random pivot from a[lo..hi].
  3. Partition a[lo..hi] with pointers lt, i, gt so that a[lo..lt-1] < pivot, a[lt..gt] == pivot, a[gt+1..hi] > pivot.
  4. If target < lt, continue in [lo, lt - 1]; if target > gt, continue in [gt + 1, hi]; otherwise return pivot.
static int optimal(int[] nums, int k) {
    int[] a = nums.clone();
    int target = a.length - k;                       // index of the answer in ascending order
    int lo = 0, hi = a.length - 1;
    Random rnd = new Random(12345);                  // fixed seed: reproducible, still a spread of pivots
    while (true) {
        int pivot = a[lo + rnd.nextInt(hi - lo + 1)];
        int lt = lo, i = lo, gt = hi;                // a[lo..lt-1] < pivot, a[lt..i-1] == pivot, a[gt+1..hi] > pivot
        while (i <= gt) {
            if (a[i] < pivot) swap(a, lt++, i++);
            else if (a[i] > pivot) swap(a, i, gt--);
            else i++;
        }
        if (target < lt) hi = lt - 1;                // answer is among the smaller values
        else if (target > gt) lo = gt + 1;           // answer is among the larger values
        else return pivot;                           // answer lands inside the block equal to pivot
    }
}

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

The loop always terminates: the equal block contains at least the pivot itself, so every round either returns or strictly shrinks [lo, hi], and the invariant lo ≤ target ≤ hi holds throughout.

Complexity. Time O(n) on average, O(n^2) in the (astronomically unlikely, with random pivots) worst case. Space O(n) for the copy, O(1) beyond it; partitioning the caller’s array directly would make it O(1) overall at the cost of reordering the input. If a guaranteed O(n) is ever required, the median-of-medians pivot rule provides it, but it is slower in practice and rarely asked for.

Edge cases and pitfalls