← All Posts
DSA · Heaps · Part 7 of 17

Top-K & the Bounded Heap Pattern

“Find the k largest” is the most-asked heap question in interviews, and the answer contains a counterintuitive move that trips up most people: to find the k largest, you use a min-heap. Once that inversion clicks, an entire family of problems collapses into one template.

Why a Min-Heap for the k Largest

Keep a heap of size exactly k holding the best candidates seen so far. When a new element arrives you must answer one question: is this better than the worst thing I am currently keeping?

So you need O(1) access to the worst element of your kept set — because that is the one you would evict. For the k largest, the worst kept element is the smallest of them. A min-heap puts exactly that at the root.

maintain a min-heap of size k
for each element x:
    if heap.size() < k:            heap.push(x)
    else if x > heap.top():        heap.pop(); heap.push(x)
// the heap now holds the k largest; heap.top() is the k-th largest

The root doubles as a rejection filter. Any element that fails x > heap.top() is discarded in a single comparison with no heap operation at all — and on most data, the overwhelming majority of elements are discarded that way. That is why this runs so much faster than its bound suggests.

The rule, in one line: k largest → min-heap of size k. k smallest → max-heap of size k. The heap is always ordered opposite to the direction you are selecting, because the root must be the candidate you are most willing to throw away.

Implementation

#include <queue>
#include <vector>
#include <functional>

std::vector<int> kLargest(const std::vector<int>& v, int k) {
    std::priority_queue<int, std::vector<int>, std::greater<int>> heap;   // MIN-heap
    for (int x : v) {
        if ((int)heap.size() < k) {
            heap.push(x);
        } else if (x > heap.top()) {
            heap.pop();
            heap.push(x);
        }
    }
    std::vector<int> out;
    while (!heap.empty()) { out.push_back(heap.top()); heap.pop(); }
    return out;                      // ascending; reverse if you want descending
}

The k-th largest element specifically is just heap.top() at the end — no extraction loop needed.

Complexity, and Why It Beats Sorting

ApproachTimeSpaceStreaming?
Sort, take last kO(n log n)O(n)No
Max-heap of all n, pop kO(n + k log n)O(n)No
Min-heap of size kO(n log k)O(k)Yes
QuickselectO(n) average, O(n2) worstO(1)No

Two properties make the bounded heap the right default:

Space is O(k), not O(n). This is often the decisive point. Finding the top 10 of a billion-element stream needs a heap of 10 elements. You cannot sort what you cannot hold, and you cannot quickselect a stream at all — quickselect needs random access to the entire array.

It is online. Elements can arrive one at a time from a socket, a log tail, or a database cursor. At every moment the heap holds the correct answer for everything seen so far. Sorting and quickselect both require the full dataset before they produce anything.

If you do have the whole array in memory and only need the k largest as an unordered set, std::nth_element (quickselect) is O(n) average and will usually win:

std::nth_element(v.begin(), v.begin() + k, v.end(), std::greater<int>());
// v[0..k) are the k largest, in unspecified order; O(n) average

Top-K Frequent Elements

The most common variant. Count first, then run the same template over the count pairs:

#include <unordered_map>
#include <queue>
#include <vector>

std::vector<int> topKFrequent(const std::vector<int>& nums, int k) {
    std::unordered_map<int, int> freq;
    for (int x : nums) ++freq[x];

    using Entry = std::pair<int, int>;             // (count, value)
    std::priority_queue<Entry, std::vector<Entry>, std::greater<Entry>> heap;

    for (const auto& kv : freq) {
        heap.push({kv.second, kv.first});
        if ((int)heap.size() > k) heap.pop();      // evict the least frequent
    }

    std::vector<int> out;
    while (!heap.empty()) { out.push_back(heap.top().second); heap.pop(); }
    return out;
}

Note the ordering of the pair: {count, value}, because std::pair compares lexicographically and we want count to dominate. Reversing the fields silently sorts by value instead — a bug that produces plausible-looking wrong answers.

Here m is the number of distinct values, so the cost is O(n + m log k). When k is close to m, bucket sort by frequency is O(n) and better; frequencies are bounded by n, so you can index buckets directly.

K Closest Points to the Origin

Same template, different key. Wanting the k smallest distances means a max-heap, so the root is the farthest point currently kept:

std::vector<std::vector<int>> kClosest(std::vector<std::vector<int>>& pts, int k) {
    // max-heap keyed by squared distance
    std::priority_queue<std::pair<long long, int>> heap;

    for (int i = 0; i < (int)pts.size(); ++i) {
        long long x = pts[i][0], y = pts[i][1];
        long long d = x * x + y * y;               // no sqrt: it is monotonic
        heap.push({d, i});
        if ((int)heap.size() > k) heap.pop();      // drop the farthest
    }

    std::vector<std::vector<int>> out;
    while (!heap.empty()) { out.push_back(pts[heap.top().second]); heap.pop(); }
    return out;
}

Two details worth stealing. Comparing squared distances avoids sqrt entirely — squaring is monotonic on non-negative values, so it preserves ordering while staying in integer arithmetic. And the accumulator is long long: with coordinates up to 104, x*x + y*y reaches 2 × 108, which fits in int, but the margin evaporates the moment the constraints grow. Overflow inside a comparator is a nasty failure because it breaks the strict weak ordering rather than producing a visibly wrong number.

The Template

Every problem in this family reduces to four decisions:

  1. Direction. Selecting the largest? Min-heap. Smallest? Max-heap.
  2. Key. What scalar is being compared — the value, a count, a distance, a ratio?
  3. Bound. Push, then pop if size() > k. Keeps the heap at exactly k.
  4. Filter. Optionally skip the push entirely when the element cannot beat top(), which avoids two heap operations per rejected element.
// generic bounded-heap skeleton
template <class T, class Compare>
std::vector<T> topK(const std::vector<T>& items, std::size_t k, Compare worst_first) {
    std::priority_queue<T, std::vector<T>, Compare> heap(worst_first);
    for (const T& x : items) {
        heap.push(x);
        if (heap.size() > k) heap.pop();
    }
    std::vector<T> out;
    while (!heap.empty()) { out.push_back(heap.top()); heap.pop(); }
    return out;
}

Practice