Top-k with a heap

Why this matters

A video site has a list of the ten most-watched videos of the day. A shop wants the five biggest orders of the hour. A teacher wants the three highest scores in a class of 40 students. Each of these asks for the top k values out of many. Here k is a small number, such as 3, 5 or 10, and the list can be much longer.

The simplest way to get the top values is to sort the whole list from largest to smallest and take the first k. That works, but it costs more work than the question needs, and it needs the whole list in memory at once. A second method uses a container called a heap. It holds only k values, and it never needs to sort everything. This lesson shows that method in C++ with std::priority_queue. It also covers the mistakes that cause a heap-based answer to come out wrong.

The same idea applies when the values come in as a stream, which means one value at a time from a source too large to store, such as a log file or a live feed. The program reads each value once, keeps the current top k, and throws the rest away.

Words you need first

A heap is a tree-shaped structure stored in a list. It is built so that one value always sits on top, and that value is the one that ranks highest under a rule you choose. The rule can be "smallest first" or "largest first". The heap does not keep the whole list in order. It only keeps the top value easy to reach.

A min-heap puts the smallest value on top. A max-heap puts the largest value on top. In C++, std::priority_queue is a heap. By default it is a max-heap, because the default rule is "less than", which puts the largest value on top.

A comparator is a rule that says which of two values should rank higher. In C++ it is a function or a function object that returns true when the first value should sit below the second one. The phrase "ranks higher" is the key. With std::greater<int>, the rule says that a smaller number ranks higher, so the smallest number sits on top. This is how you get a min-heap.

A top-k answer is the set of the k largest values. For the list 42 7 19 88 3 56 21 90 14 with k equal to 3, the top-3 values are 90, 88 and 56.

The heap operations you use are these. push(x) adds a value. top() reads the value on top without removing it. pop() removes the top value and returns nothing, so read top() first. size() and empty() tell you how many values are in the heap. The cost of push and pop is O(log k), which means it grows slowly with k. For k equal to 1000, log2(k) is about 10. The cost of top() is O(1).

The slow way first: sort everything

Start with the method everyone knows. Sort the list from largest to smallest and read the first k values.

#include <algorithm>
#include <cstddef>
#include <functional>
#include <iostream>
#include <vector>
 
int main() {
    std::vector<int> scores{42, 7, 19, 88, 3, 56, 21, 90, 14};
    std::size_t k = 3;
 
    std::sort(scores.begin(), scores.end(), std::greater<int>());  // largest first
    for (std::size_t i = 0; i < k; ++i) {
        std::cout << scores[i] << ' ';
    }
    std::cout << '\n';
}

Output:

90 88 56 

This code is short and easy to check, and for a list of nine values it does not matter which method you use. For a list of a million values, sorting does about n times log n comparisons, which is about 20 million for a million values. It also keeps all one million values in memory. The question only asks for three of them, so most of the sorted work is wasted.

The heap method does about n times log k work. With k equal to 3, that is about two comparisons per value. It also needs room for only k values, no matter how long the list is.

The idea: a waiting room with k places

Picture a waiting room with k chairs. Each chair holds one value. When the room is not full, every new value takes a chair. When the room is full, a new value can only get in if it is bigger than the weakest value in the room. The weakest value, the one with the smallest number, is the one the heap has on top. If the new value wins, the weakest value leaves and the new value takes its chair. If it does not win, it is thrown away.

When the whole list has been read, the room holds the top k values. A value is thrown away only when the room already holds k values that are at least as big, so no value that belongs in the top k is ever lost. Here is the trace for the list 42 7 19 88 3 56 21 90 14 with k equal to 3. The heap is a min-heap, so the smallest value is written first.

The last row is the answer: 56, 88 and 90. The printed order in the program is 90, 88, 56, which is the reverse of the order the heap gives them in.

Worked example 1: the top three of a short list

Here is the same trace as code. The heap is declared with std::greater<int>, which makes it a min-heap.

#include <algorithm>
#include <cstddef>
#include <functional>
#include <iostream>
#include <queue>
#include <vector>
 
int main() {
    std::vector<int> scores{42, 7, 19, 88, 3, 56, 21, 90, 14};
    std::size_t k = 3;
 
    // Min-heap: the smallest value sits on top.
    std::priority_queue<int, std::vector<int>, std::greater<int>> heap;
 
    for (int x : scores) {
        if (heap.size() < k) {
            heap.push(x);
        } else if (x > heap.top()) {
            heap.pop();
            heap.push(x);
        }
    }
 
    std::vector<int> result;
    while (!heap.empty()) {
        result.push_back(heap.top());
        heap.pop();
    }
    std::reverse(result.begin(), result.end());
 
    for (int x : result) {
        std::cout << x << ' ';
    }
    std::cout << '\n';
}

Output:

90 88 56 

The type std::priority_queue<int, std::vector<int>, std::greater<int>> has three parts. The first is the type of value stored, int. The second is the list the heap uses for storage, here std::vector<int>. The third is the comparator, here std::greater<int>. The second and third parts can be left out when you want the default: std::priority_queue<int> uses a std::vector<int> and std::less<int>.

The loop has two branches. While the heap holds fewer than k values, every new value is pushed. After that, the new value is compared with heap.top(), the weakest value in the room. If the new value is bigger, the weakest value is removed with pop() and the new value is pushed. This is the same rule as the trace.

The printing step looks odd at first. The heap gives values smallest first, so the first top() read is the weakest of the three. The code collects them in that order, then reverses the list, so the answer is printed largest first.

Worked example 2: the top five of a stream

The next program does not keep the list at all. It makes one million values, one at a time, and each value is looked at once and then forgotten, except for the top five. A real program would read the values from a file or a network connection, but the logic is the same.

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <iostream>
#include <queue>
#include <random>
#include <vector>
 
int main() {
    const std::size_t k = 5;
    std::mt19937 gen(7);  // a fixed seed, so the run is repeatable
 
    std::priority_queue<std::uint32_t, std::vector<std::uint32_t>, std::greater<>> heap;
    std::size_t largest_size = 0;
 
    for (int i = 0; i < 1'000'000; ++i) {
        std::uint32_t x = gen() % 1'000'000;  // one value arrives at a time
        if (heap.size() < k) {
            heap.push(x);
        } else if (x > heap.top()) {
            heap.pop();
            heap.push(x);
        }
        largest_size = std::max(largest_size, heap.size());
    }
 
    std::vector<std::uint32_t> result;
    while (!heap.empty()) {
        result.push_back(heap.top());
        heap.pop();
    }
    std::reverse(result.begin(), result.end());
 
    for (std::uint32_t x : result) {
        std::cout << x << ' ';
    }
    std::cout << "\nlargest number of values held at once: " << largest_size << '\n';
}

Output:

999998 999995 999992 999991 999989 
largest number of values held at once: 5

The second line of output is the point of the example. The program read one million values but never held more than five at once. The variable largest_size tracks the biggest size the heap reached, and it stays at 5.

Two details are new here. The number 1'000'000 uses a single quote as a digit separator. C++14 and later allow it, and it makes long numbers easier to read. The comparator is std::greater<> with empty angle brackets. This is a version that works for any type that supports >, so the code does not need to repeat the type name.

The fixed seed 7 makes the random generator produce the same numbers on every run. The random values are not the point of the example. The point is that a heap of size five can find the top five of any stream, however long it is.

Worked example 3: top k records with a custom rule

Values are often part of a bigger record. A player has a name and a score, and the top three players are the three with the highest scores. The heap needs a rule that looks at the score only. This program uses a lambda, which is a small unnamed function written inside the code.

#include <algorithm>
#include <cstddef>
#include <iostream>
#include <queue>
#include <string>
#include <vector>
 
struct Player {
    std::string name;
    int score;
};
 
int main() {
    std::vector<Player> players{
        {"Asha", 310}, {"Ben", 275}, {"Chen", 402},
        {"Dara", 355}, {"Eli", 290}, {"Fatima", 402},
    };
    std::size_t k = 3;
 
    // The comparator returns true when a has the higher score. A true result
    // means a sits below b, so the lowest score ends up on top.
    auto lower_score_first = [](const Player& a, const Player& b) {
        return a.score > b.score;
    };
    std::priority_queue<Player, std::vector<Player>, decltype(lower_score_first)>
        heap(lower_score_first);
 
    for (const Player& p : players) {
        if (heap.size() < k) {
            heap.push(p);
        } else if (p.score > heap.top().score) {
            heap.pop();
            heap.push(p);
        }
    }
 
    std::vector<Player> best;
    while (!heap.empty()) {
        best.push_back(heap.top());
        heap.pop();
    }
    std::reverse(best.begin(), best.end());
 
    for (const Player& p : best) {
        std::cout << p.name << ' ' << p.score << '\n';
    }
}

Output:

Fatima 402
Chen 402
Dara 355

The comparator in this program looks backwards at first. It returns a.score > b.score, so it returns true when a has the higher score. A true result means a sits below b in the heap, so the lowest score ends up on top, which is the min-heap rule. The lambda is passed to the constructor, because the heap has to be told which rule to use. The type decltype(lower_score_first) gives the compiler the lambda's type.

Two players have 402 points. The program keeps both of them. Chen took a free chair, and Fatima beat the weakest value in the room, which was Asha's 310, when Fatima arrived. Their order in the output is not something to rely on. When two values are equal, a heap does not promise which one comes out first. If the order matters, add a second rule to the comparator, such as the name, so that every pair has a clear winner.

Common mistakes

Forgetting the comparator and getting a max-heap. The default std::priority_queue<int> puts the largest value on top. The code in worked example 1 then compares each new value with the largest value in the room, not the weakest. It replaces the largest value each time a bigger one arrives, which is the reverse of what the lesson needs. This program uses the default heap with the same loop:

#include <algorithm>
#include <cstddef>
#include <iostream>
#include <queue>
#include <vector>
 
int main() {
    std::vector<int> scores{42, 7, 19, 88, 3, 56, 21, 90, 14};
    std::size_t k = 3;
 
    // Mistake: no comparator, so this is a max-heap. The largest value sits on top.
    std::priority_queue<int> heap;
 
    for (int x : scores) {
        if (heap.size() < k) {
            heap.push(x);
        } else if (x > heap.top()) {
            heap.pop();
            heap.push(x);
        }
    }
 
    std::vector<int> result;
    while (!heap.empty()) {
        result.push_back(heap.top());
        heap.pop();
    }
    std::reverse(result.begin(), result.end());
    for (int x : result) {
        std::cout << x << ' ';
    }
    std::cout << '\n';
}

Output:

7 19 90 

The correct answer is 90 88 56. The wrong program prints 7 19 90, which is three values from the list, but not the three largest. Each time a bigger value arrives, the heap throws out its biggest value, so the biggest values are the ones that leave.

Taking the top of an empty heap when k is 0. If k is 0, the first branch never runs, because heap.size() is never less than 0. The code then reads heap.top() on an empty heap. The program below shows that case:

#include <cstddef>
#include <functional>
#include <iostream>
#include <queue>
#include <vector>
 
int main() {
    std::vector<int> scores{42, 7, 19};
    std::size_t k = 0;
    std::priority_queue<int, std::vector<int>, std::greater<int>> heap;
 
    for (int x : scores) {
        if (heap.size() < k) {
            heap.push(x);
        } else if (x > heap.top()) {  // heap is empty here when k is 0
            heap.pop();
            heap.push(x);
        }
    }
    std::cout << "done, size " << heap.size() << '\n';
}

Built with g++ -std=c++20 -Wall -D_GLIBCXX_ASSERTIONS, the flag that turns on the library's checks, the program stops with this message:

/usr/include/c++/15/bits/stl_queue.h:825: std::priority_queue<_Tp, _Sequence, _Compare>::const_reference std::priority_queue<_Tp, _Sequence, _Compare>::top() const [with _Tp = int; _Sequence = std::vector<int>; _Compare = std::greater<int>; const_reference = const int&]: Assertion '!this->empty()' failed.

Without the checks, reading top() on an empty heap is undefined behaviour, so the program may print garbage or crash somewhere else. The fix is one line at the top of the function: if (k == 0) return {}; for a function that returns a list.

Comparing a size with a signed number. The size() method returns an unsigned type. Writing int k and then heap.size() < k compares an unsigned number with a signed one. The program below does that:

#include <iostream>
#include <queue>
 
int main() {
    std::priority_queue<int> heap;
    int k = 3;
    heap.push(4);
    if (heap.size() < k) {
        std::cout << "room\n";
    }
}

Compiled with g++ -std=c++20 -Wall, it prints this warning. The program still builds, and this small case still behaves, but a negative k would be read as a huge unsigned number.

warning: comparison of integer expressions of different signedness: ‘std::priority_queue<int>::size_type’ {aka ‘long unsigned int’} and ‘int’ [-Wsign-compare]

The fix is to declare k as std::size_t, as the examples do.

Replacing the top with the wrong test. The replacement test must compare the new value with the weakest value in the room, which is the heap's top when the rule is min-heap. In the players program, the test reads p.score > heap.top().score. Dropping .score, or writing heap.top() > p on a min-heap, changes the meaning of the test. Read the test aloud as "does the new value beat the weakest value in the room?" and check that the words match the code.

Expecting an order among equal values. As worked example 3 showed, two values that compare as equal can come out in either order. If the output must be stable, add a tie-breaker to the comparator.

Practice

Exercise 1. The list is 5 9 1 7 and k is 2. Use the waiting room method to find the top two values, and list the heap after each step.

Show the answer

Value 5 takes a chair, and value 9 takes a chair. The heap is 5, 9. Value 1 does not beat 5, so it is thrown away. Value 7 beats 5, so 5 leaves and 7 takes its chair. The heap is 7, 9. The top two values are 9 and 7.

Exercise 2. Change worked example 1 so it returns the three smallest values, not the three largest. Give the change, and say what the output is for the same list.

Show the answer

Declare the heap as std::priority_queue<int>, which is the default max-heap, and change the test to x < heap.top(). The max-heap keeps the largest of the kept values on top, so each new smaller value replaces it. The output for the list 42 7 19 88 3 56 21 90 14 is 3 7 14. Running that version gives exactly this output.

Exercise 3. The lesson says x > heap.top() and not x >= heap.top(). For plain integers, does the output change? For players, why might it change?

Show the answer

For plain integers, the output does not change, because two equal integers look the same and it does not matter which one stays. For players, a tie in score means the two records differ in their names. With >, the value already in the heap stays, and the new record is thrown away. With >=, a new record that ties with the weakest record in the room replaces it, so the name in the output can change.

Exercise 4. A list has n values, and you keep the top k. Explain why the total cost is O(n log k) and why sorting costs O(n log n).

Show the answer

Each value is pushed at most once and popped at most once. A push or a pop on a heap of size k takes about log k steps. So the total is about n times log k. Sorting compares about n times log n pairs, because the list has n values to arrange. When k is much smaller than n, log k is much smaller than log n, which is why the heap method wins.

Exercise 5. Write one line of code that makes k equal to 0 return an empty result from a function, and explain why that line must come before the loop.

Show the answer

Write if (k == 0) { return {}; } at the top of the function, before the loop. Without it, the loop reads heap.top() on an empty heap when the first value arrives, which is the mistake shown in the common mistakes section.

Summary

• The top-k problem asks for the k largest values, and a heap of size k answers it without sorting the whole list.
• A min-heap keeps the smallest of the kept values on top, so each new value is compared with the weakest value in the room.
• A new value replaces the top only when it is bigger, which is the test x > heap.top().
• The heap holds at most k values, so the same code works on a stream that is too large to store.
• Forgetting the comparator gives a max-heap, which keeps the wrong values, and k equal to 0 reads an empty heap.
• Each value costs about log k work, so the whole run takes O(n log k) time.