Skip to main content

Top K and the Kth Element

Editorial

The root is the bar to clear

To find the 3 largest of a million values you don't need the million in order. You need the best 3 so far, and a quick way to tell whether a new value beats any of them. Keep those 3 in a min-heap: its root is the weakest of the best, the 3rd largest so far. A new value bigger than the root takes its place; anything else can never be in the top 3. Each value costs one look at the root and, at most, one replace, O(log k):

A zone titled the stream, one value at a time holds eight shaded columns, one per value of 4 9 1 7 3 8 2 6, each as tall as its value. The first three, 4, 9 and 1, are purple, room: the heap of three fills. From then on a dashed orange line at the root's value runs across the columns, rising from 1 to 4 to 7, with a pill reading the root: 7 and 8 rise above it in green and get in, 3, 2 and 6 stay under it, pale, skipped. Under each column a pill says room, or the value against the root, such as > 1 or < 4. The line runs on into a zone titled the final heap, to its root, 7, ringed in orange over 9 and 8, with a pill reading 7 is the 3rd largest. Focal line: a value gets in only by clearing the root. Title: the root is the bar to clear.

The kth largest with a heap of size k

For the k largest, a min-heap of at most k values; for the k smallest, the mirror, a max-heap:

// The k largest so far in a min-heap: its top is the kth largest, the bar a new value has to clear.
int findKthLargest(vector<int>& nums, int k) {
    priority_queue<int, vector<int>, greater<int>> heap;
    for (int x : nums) {
        if ((int)heap.size() < k) heap.push(x);
        else if (x > heap.top()) { heap.pop(); heap.push(x); }   // x beats the weakest of the best
    }
    return heap.top();
}
import heapq


def find_kth_largest(nums, k):
    """The k largest so far in a min-heap: heap[0] is the kth largest, the bar a new value has to clear."""
    heap = []
    for x in nums:
        if len(heap) < k:
            heapq.heappush(heap, x)
        elif x > heap[0]:
            heapq.heapreplace(heap, x)       # x beats the weakest of the best and takes its place
    return heap[0]
// The k largest so far in a min-heap: peek() is the kth largest, the bar a new value has to clear.
static int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> heap = new PriorityQueue<>();
    for (int x : nums) {
        if (heap.size() < k) heap.offer(x);
        else if (x > heap.peek()) { heap.poll(); heap.offer(x); }   // x beats the weakest of the best
    }
    return heap.peek();
}
// The k largest so far in Introduction to Heaps' MinHeap: a[0] is the kth largest, the bar to clear.
function findKthLargest(nums, k) {
    const heap = new MinHeap();
    for (const x of nums) {
        if (heap.a.length < k) heap.push(x);
        else if (x > heap.a[0]) { heap.pop(); heap.push(x); }   // x beats the weakest of the best
    }
    return heap.a[0];
}

Each printed 5 4 1: Kth Largest Element in an Array's two examples and [1] with k = 1. That's O(n log k) time and O(k) memory. Step through it on a stream:

The 3 largest of [4, 9, 1, 7, 3, 8, 2, 6] — The stream is the row below; the tree is the heap of the best so far. Green is a value that just got in; amber, the root a value failed to beat.

Below, you keep the best k. Each value arrives; push it while the heap has room, replace the root when the value beats it, and skip it otherwise. A wrong choice is refused, and the root is always the kth largest so far.

Keep the Best K: push, replace the root, or skip

Against sorting and quickselect

Two other ways get the kth largest. Sort everything and index: O(n log n), and it needs all n at once. Or quickselect: partition around a pivot as quicksort does, then carry on in the one part that holds index n − k, O(n) on average (Quick Sort and Quickselect builds it); C++ ships one as nth_element. Timed in C++ on the same 1,000,000 random values, the best of 5 runs for each k:

kHeap of size kSort, then indexQuickselect
101.3 ms77.3 ms5.7 ms
1,0001.9 ms77.3 ms5.7 ms
10,0006.4 ms77.1 ms5.7 ms
100,00033.2 ms77.1 ms5.8 ms
500,00065.1 ms77.1 ms9.3 ms

At k = 10 and 1,000 the heap wins: nearly every value loses to the root after one comparison. As k grows, more values get in and each costs O(log k): the two are about even at 10,000, quickselect is well ahead by 100,000, and at 500,000 the heap is nearly as slow as sorting. Sorting costs the same at every k, because it orders all n. The heap keeps one advantage at any k: it never needs the values at once, which is the whole point on a stream.

On a stream, on counts, on distances

Kth Largest Element in a Stream is the same heap kept between calls:

import heapq


class KthLargest:
    def __init__(self, k, nums):
        self.k, self.heap = k, []
        for x in nums:
            self.add(x)

    def add(self, val):
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, val)
        elif val > self.heap[0]:
            heapq.heapreplace(self.heap, val)
        return self.heap[0]

With k = 3 and [4, 5, 8, 2], adding 3, 5, 10, 9 and 4 returned [4, 5, 5, 8, 8], the problem's example.

Top K Frequent Elements and K Closest Points to Origin first compute a key, a count or a distance, and then keep the best k by that key. For the closest points the "best" are the smallest distances, so the heap is a max-heap of size k, here with the distances negated:

from collections import Counter
import heapq


def top_k_frequent(nums, k):
    heap = []                                # (count, value): the k most frequent so far, the least of them on top
    for value, count in Counter(nums).items():
        if len(heap) < k:
            heapq.heappush(heap, (count, value))
        elif count > heap[0][0]:
            heapq.heapreplace(heap, (count, value))
    return sorted(v for _, v in heap)
import heapq


def k_closest(points, k):
    heap = []                                # (-distance², x, y): the farthest of the k closest on top
    for x, y in points:
        d = -(x * x + y * y)
        if len(heap) < k:
            heapq.heappush(heap, (d, x, y))
        elif d > heap[0][0]:
            heapq.heapreplace(heap, (d, x, y))
    return sorted([x, y] for _, x, y in heap)

The first returned [1, 2] for [1, 1, 1, 2, 2, 3] with k = 2. The second returned [[-2, 2]] and [[-2, 4], [3, 3]] for its problem's two examples.

A nearly sorted array, and the best pair sums

Sort Nearly Sorted Array: every value is at most k places from where it belongs. Then the smallest of the first k + 1 values must be first, and after it the smallest of the next window, and so on. A min-heap of k + 1 values always holds the next one:

Three rows, one per value sorted out of 6 5 3 2 8 10 9, where each value is at most 3 places from its sorted spot, across three zones: read the next value, a heap of k + 1, its top is next. In each row the values just read are purple tiles, first 6 5 3 2, then 8, then 10, the rest pale; beside them, the min-heap of 4 they leave, 2 over 5 and 3, then 3 over 5 and 8, then 5 over 6 and 8, its top green; and an arrow takes the top into the output, which grows 2, then 2 3, then 2 3 5, the newest ringed. Focal line: the next value is always among the k + 1 the heap holds. Title: the next value is within reach.
import heapq


def sort_nearly_sorted(a, k):
    """Every value is at most k places from where it belongs."""
    heap = a[:k + 1]
    heapq.heapify(heap)
    out = []
    for x in a[k + 1:]:
        out.append(heapq.heapreplace(heap, x))   # the smallest of these k + 1 is next; x takes its place
    while heap:
        out.append(heapq.heappop(heap))
    return out

It sorted [6, 5, 3, 2, 8, 10, 9] with k = 3 into [2, 3, 5, 6, 8, 9, 10], in O(n log k).

Maximum Sum Combination asks for the k largest sums of one value from each of two arrays. Sort both from the largest. The best sum is the two first values; every next best is one step from a sum already taken, down one array or the other. So a heap of those neighbour candidates gives the sums in order:

A zone titled every a + b holds a 4 by 4 grid of sums: rows a sorted from the largest, 4 3 2 1, on teal chips down the side, columns b sorted the same way, 6 5 2 1, on purple chips across the top. After 4 pops the taken cells are green and numbered in order: 10 in the corner, the 9 to its right, the 9 below it, then the 8. The cells right of and below them, 8, 7, 6 and 5, are orange, the heap's candidates; the other 8 cells are pale, never built. A zone titled after 4 pops lists taken, best first, 10 9 9 8; in the heap, next in line, 8 7 6 5; never built, 8 of the 16 sums. Focal line: every next best sum is next door to one already taken. Title: the next best sum is next door.
import heapq


def max_sum_combinations(a, b, k):
    a, b = sorted(a, reverse=True), sorted(b, reverse=True)
    heap, seen, out = [(-(a[0] + b[0]), 0, 0)], {(0, 0)}, []
    while len(out) < k:
        s, i, j = heapq.heappop(heap)
        out.append(-s)
        for ni, nj in ((i + 1, j), (i, j + 1)):     # the next best sums sit next to one already taken
            if ni < len(a) and nj < len(b) and (ni, nj) not in seen:
                seen.add((ni, nj))
                heapq.heappush(heap, (-(a[ni] + b[nj]), ni, nj))
    return out

It returned [10, 9, 9, 8] for [1, 4, 2, 3] and [2, 5, 1, 6] with k = 4, without building all 16 sums: O(k log k) after the two sorts.

Spotting it, and when it breaks

The clues. "The kth largest" or "kth smallest", "the top k", "the k closest", "the k most frequent", and anything about a stream where you can't keep everything: a heap of size k.

The near-misses. k close to n: flip it, since the kth largest is the (n − k + 1)th smallest, and keep a max-heap of n − k + 1. k near n / 2 on an array you already hold: quickselect, as timed above. Every value in order, not just k: sort. The single largest: one pass. The kth of a sorted matrix or of merged sorted lists: a heap of candidates, as in the pair sums, not a heap of everything. A sliding window's maximum: a monotonic deque (Monotonic Deque: the Window's Maximum). And counts are whole numbers up to n, so Top K Frequent Elements can also put each value in a bucket by its count and read the buckets from the top, O(n).

When it breaks. Choosing the wrong heap flips the answer: the k largest need the min-heap. And a heap of every value works but pays O(n log n) time and O(n) memory, which is what the size-k heap avoids.

Bugs that cost a submission

A max-heap for the k largest.

import heapq


def find_kth_largest(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, -x)                 # a max-heap of size k
        if len(heap) > k:
            heapq.heappop(heap)                  # throws out the largest
    return -heap[0]

For [3, 2, 1, 5, 6, 4] with k = 2 it returned 2; the answer is 5. Popping the top of a max-heap throws out the largest, so it keeps the k smallest. The k largest need a min-heap, whose top is the one to throw out.

Never cutting the heap back to k.

import heapq


def find_kth_largest(nums, k):
    heap = []
    for x in nums:
        if not heap or x > heap[0]:
            heapq.heappush(heap, x)              # never cut back to k
    return heap[0]

It returned 3: the heap kept every value that ever beat the root, so its top is no longer the kth largest. Push while there's room, then replace.

Re-sorting a stream on every add.

class KthLargest:
    def __init__(self, k, nums):
        self.k, self.values = k, list(nums)

    def add(self, val):
        self.values.append(val)
        self.values.sort()                       # every value, re-sorted on every add
        return self.values[-self.k]

For 20,000 adds with k = 10 it took 1,094.7 ms, against 5.9 ms for the heap of size k, with the same 20,000 answers. The heap keeps 10 values; the sort keeps, and re-sorts, every one.

Check yourself

Quiz

1 / 6

Multiple choice

To keep the k largest values seen so far, which heap of size k do you use?

Practice problems

In this order:

  1. Kth Largest Element in a Stream: the heap kept between calls.
  2. Kth Largest Element in an Array: a min-heap of size k, or quickselect.
  3. Kth Smallest Element: the mirror, a max-heap of size k.
  4. Top K Frequent Elements: count, then the size-k heap by count.
  5. K Closest Points to Origin: the size-k heap by distance.
  6. Sort Nearly Sorted Array: a window of k + 1.
  7. Maximum Sum Combination: a heap of neighbour candidates.

Next, Merging and Scheduling with a Heap: one candidate from each list, and the next task to run.