Heap Kth-Element Selection

Description

"What's the 10th largest transaction today?" Sorting all n values costs O(n log n) and holds everything in order just to read one position. When k is small, a heap that remembers only the best k is much cheaper - and for a one-off query on an array, quickselect is cheaper still.

This skill covers the two standard tools: a min-heap capped at k, which works on streams and leaves the input alone, and quickselect, which finds one rank in linear expected time by rearranging an array. It also covers picking between them, and the comparator and index mistakes that make either one quietly wrong.

We'll go from simple to more technical:

1. The pattern: a min-heap of the k largest so far, evicting the smallest when it grows past k.
2. A second flavor: quickselect, which partitions and keeps only the side holding the answer.
3. Where top-k shows up (slowest requests, largest orders, ORDER BY ... LIMIT).
4. What breaks it: the wrong heap, no size cap, the n - k index, fixed pivots.

The approach

int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> heap = new PriorityQueue<>();     // MIN-heap of the k largest so far
    for (int x : nums) {
        heap.offer(x);
        if (heap.size() > k) heap.poll();                     // evict the smallest of the kept ones
    }
    return heap.peek();                                       // the smallest of the k largest
}

import heapq


def find_kth_largest(nums: list[int], k: int) -> int:
    heap = []                                  # MIN-heap of the k largest so far (heapq: functions on a list)
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)                # evict the smallest of the kept ones
    return heap[0]                             # the smallest of the k largest

# the standard library does this in one call: heapq.nlargest(k, nums)[-1]

The heap holds the k largest values seen so far, with the smallest of them on top - exactly the one to evict when something bigger arrives. A value that gets evicted has k values at least as large already kept, so it can't be among the k largest. It's a min-heap even though we want the largest, because the heap has to expose the weakest member. Cost: O(n log k) time, O(k) space. When k is close to n, flip it: the k-th largest is the (n - k + 1)-th smallest, so keep that many values in a max-heap. (Python's heapq functions are a min-heap: push negated values, or on Python 3.14+ use heapq.heappush_max and heappop_max.)

A second flavor: quickselect

For an array you can rearrange, partition it as quicksort does and recurse only into the side that holds the answer:

int kthLargest(int[] a, int k) {
    int target = a.length - k;                  // the k-th largest sits here in ascending order
    int lo = 0, hi = a.length - 1;
    Random rnd = new Random();
    while (true) {
        int p = partition(a, lo, hi, lo + rnd.nextInt(hi - lo + 1));
        if (p == target) return a[p];
        if (p < target) lo = p + 1; else hi = p - 1;
    }
}

int partition(int[] a, int lo, int hi, int pivotIndex) {   // Lomuto: smaller values to the left
    swap(a, pivotIndex, hi);
    int store = lo;
    for (int i = lo; i < hi; i++) if (a[i] < a[hi]) swap(a, i, store++);
    swap(a, store, hi);
    return store;
}

import random


def kth_largest(a: list[int], k: int) -> int:
    target = len(a) - k                        # the k-th largest sits here in ascending order
    lo, hi = 0, len(a) - 1
    while True:
        p = partition(a, lo, hi, random.randint(lo, hi))
        if p == target:
            return a[p]
        if p < target:
            lo = p + 1
        else:
            hi = p - 1


def partition(a: list[int], lo: int, hi: int, pivot_index: int) -> int:   # Lomuto: smaller values to the left
    a[pivot_index], a[hi] = a[hi], a[pivot_index]
    store = lo
    for i in range(lo, hi):
        if a[i] < a[hi]:
            a[i], a[store] = a[store], a[i]
            store += 1
    a[store], a[hi] = a[hi], a[store]
    return store

Each round discards one side, so the expected work is n + n/2 + n/4 + ... = O(n) with a random pivot. With a fixed first-element pivot, sorted input discards one element per round: 12.5 million comparisons for 5,000 sorted values versus about 10,000 shuffled. Quickselect also reorders the caller's array.

Where you'll see it

"Top 10 slowest requests in the last hour" and "5 largest orders today" over streams too large to sort are size-k heaps. Monitoring percentiles (p99 latency) over a batch is a selection problem too; databases answer ORDER BY ... LIMIT k with a bounded heap ("top-N sort") instead of sorting every row.

Which tool for which job?

Common mistakes

Check yourself

Try to answer before opening each one - recalling beats re-reading.

1. Why does finding the k largest use a MIN-heap?

2. What does quickselect do, and why is it expected O(n)?

3. When is a size-k heap better than quickselect?

Key takeaways

• k largest: min-heap capped at k; the root is the k-th largest. O(n log k).
• For k near n, keep the n - k + 1 smallest in a max-heap.
• Quickselect: O(n) expected with a random pivot, O(n^2) with a fixed one on sorted input; it mutates the array.
• The k-th largest is index n - k in ascending order.