Topics › Heaps › Medium Problems

K-th largest element in an unsorted array.

Medium TUF ID 2839 SDE

Problem

Given an unsorted integer array nums and an integer k, return the k-th largest element. This means the element that would sit at position k if the array were sorted in descending order, so duplicates are counted separately: in [3, 3, 1] the 2nd largest is 3, not 1. The array itself should not be changed.

Constraints (LeetCode): 1 ≤ k ≤ n ≤ 10^5, -10^4 ≤ nums[i] ≤ 10^4. The follow-up asks for a solution that does not sort.

Example 1

nums = [3, 2, 1, 5, 6, 4], k = 2
answer = 5          // 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

A useful translation: the k-th largest of n elements is the element at index n - k of the array sorted ascending. All three approaches below aim at that index.

Approach 1: Brute force, sort a copy

Intuition. Sorting puts every element at its final rank, after which the answer is a single array lookup. We sort a clone so the caller’s array is left untouched. This is correct and short, but it computes the rank of every element when we only need one.

Algorithm.

  1. Clone nums and sort the clone ascending.
  2. Return the element at index 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 (plus the sort’s own small overhead).

Approach 2: Better, a min-heap of the k largest

Intuition. Scan the array while keeping the k largest values seen so far. Among those k, the smallest one is the current k-th largest. A min-heap keeps exactly that smallest value at its root. Whenever the heap grows to k + 1 elements, the root is smaller than k other values, so it can never be the answer and we discard it. At the end the heap holds the true top k, and its root is the answer.

This approach shines when k is small compared to n, or when the data arrives as a stream and cannot be stored.

Algorithm.

  1. For each x in nums: push x; if the heap size exceeds k, pop the root.
  2. Return the root.
static int better(int[] nums, int k) {
    PriorityQueue<Integer> heap = new PriorityQueue<>();    // min-heap: root = smallest of the kept k
    for (int x : nums) {
        heap.offer(x);
        if (heap.size() > k) heap.poll();                   // throw away the smallest, it cannot be the answer
    }
    return heap.peek();
}

Complexity. Time O(n log k): each element costs at most one push and one pop on a heap of size k + 1. Space O(k).

Approach 3: Optimal, quickselect

Intuition. Quicksort works by partitioning around a pivot: afterwards the pivot sits at its final sorted position, everything smaller is to its left, and everything larger is to its right. Quicksort then recurses into both sides. But we only care about one position, index n - k. After one partition we know which side contains that index, so we can discard the other side completely. That is quickselect.

With a random pivot, each partition shrinks the range to about half on average, so the total work is around n + n/2 + n/4 + ... ≈ 2n, which is linear. Sorted input cannot trick a random pivot. The seed is fixed only to make test runs reproducible; it still behaves like random choice for any input that was not built against this exact seed.

We partition three ways (Dijkstra’s Dutch national flag) instead of two. While scanning a[lo..hi] we keep:

When the scan finishes, every copy of the pivot occupies lt..gt, and those are exactly the final sorted positions of all the pivot copies. If the target index lies in lt..gt, the pivot is the answer. This matters for duplicates: with a two-way partition, an array of 10^5 equal values shrinks by only one element per round and degrades to O(n^2), whereas three-way partitioning finishes it in a single round.

Algorithm.

  1. Clone the array. Set target = n - k, lo = 0, hi = n - 1.
  2. Pick a random pivot value from a[lo..hi] and three-way partition the range, giving lt and gt.
    • When a[i] < pivot, swap it into the < block and advance both lt and i.
    • When a[i] > pivot, swap it with a[gt] and decrement gt. Do not advance i, because the element swapped in has not been examined.
    • Otherwise advance i.
  3. If target < lt, continue with hi = lt - 1. If target > gt, continue with lo = gt + 1. Otherwise return the pivot.
static int optimal(int[] nums, int k) {
    int[] a = nums.clone();                                 // drop the clone if the caller allows mutation
    int target = a.length - k;                              // index of the answer in ascending order
    int lo = 0, hi = a.length - 1;
    Random rng = new Random(2839);                          // fixed seed: reproducible, still breaks sorted input
    while (true) {
        int pivot = a[lo + rng.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;                                  // target landed in 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;
}

Dry run on Example 1 (target = 6 - 2 = 4). The pivots shown are one possible sequence of random choices:

round range pivot range after partition lt..gt decision
1 0..5 4 [3, 2, 1, 4, 6, 5] 3..3 4 > 3, so lo = 4
2 4..5 6 [5, 6] 5..5 4 < 5, so hi = 4
3 4..4 5 [5] 4..4 4 is inside: return 5

Complexity. Expected time O(n). The worst case is O(n^2) if the random pivots are consistently terrible, which is astronomically unlikely. (The median-of-medians pivot rule guarantees O(n) worst case, but its constant factor is large and it is rarely expected in interviews.) The loop is iterative, so there is no recursion stack. Space O(n) for the clone; if the caller allows the array to be reordered, drop the clone and the extra space becomes O(1).

Edge cases and pitfalls