Top K elements and merging K sorted lists

Why this matters

Many real questions are not about the whole collection. They ask for a small part of it. Which five products sold best this week? Which ten search results are most relevant out of a million? Which three words appear most in a book? Which record is the twentieth largest in a stream that never ends? A sorting method can answer each of these questions, but it does far more work than needed, because it puts everything in order just to read a few values.

A heap lets you keep only the part you need. This lesson shows the pattern that solves most "top K" problems, and a second classic problem, merging K sorted lists, which uses a heap to pick the next smallest value across many lists at once. Both problems appear often in interviews. Both are short to write once you see the pattern, and both are easy to get subtly wrong, so we will go slowly and trace every step.

You should already know push, pop and top from the earlier lessons, and you should know how to write a comparison for a priority_queue. If you are unsure about the comparison, reread the section on custom orders in the last lesson first.

The first principle: keep only what you need

Suppose you have n numbers and you want the k largest, where k is small, say 3. Here are three ways to do it.

The first way sorts all the numbers from smallest to largest and reads the last three. Sorting costs O(n log n) time and needs room for all n numbers. For a million numbers, that is a lot of work to find three.

The second way puts all the numbers in a max-heap and pops three times. Building the heap costs O(n), and each pop costs O(log n). The total is O(n + k log n), which is better, but the heap still holds all n numbers in memory.

The third way is the one that matters. Keep a min-heap of at most k numbers. For each new number, push it. If the heap now holds more than k numbers, pop the smallest one. At the end, the heap holds exactly the k largest numbers, and its smallest value is the kth largest. The cost is O(n log k) time, which is fast when k is small, and the memory is only O(k).

The third way seems backwards at first. You use a min-heap to find the largest values. The reason is that the root of a min-heap is the weakest of the values you are keeping. When a new value arrives, you only need to compare it with that weakest value. If the new value is bigger, it replaces the weakest one. If it is smaller, it is not in the top k, and you throw it away.

An everyday picture

Think of a nightclub with room for exactly k guests. The bouncer keeps the best k people inside. The bouncer also knows who is the least impressive guest inside, because that person stands nearest the door. When a new, more impressive person arrives and the room is full, the bouncer sends the least impressive guest out. A person who is less impressive than the weakest guest inside is turned away at once. The door is the root of a min-heap, and the least impressive guest is the root.

The bouncer never needs to see the whole crowd in order. The bouncer only needs to know the weakest guest inside. That is exactly what a min-heap gives in O(1).

The rule to remember

• To keep the k largest values, use a min-heap of size k.
• To keep the k smallest values, use a max-heap of size k.

The heap is the opposite of the target because you want to be able to throw out the value that is least useful, and that is the root of the opposite heap.

Top K largest: the program

The first program reads a count n, a number k, and then n values. It prints the k largest values in descending order. The input format is one line with n and k, then one line with the values.

Type this into the Input box:

8 3
5 1 9 3 7 9 2 6

The output is:

top 3: 9 9 7

The heap pops in ascending order, so the top vector is ascending. The loop that prints goes backwards to show the largest value first. Notice that 9 appears twice in the input, and it is correct to print it twice, because both copies belong in the top three.

Trace the first few values. With k = 3:

1. Push 5. Size is 1, which is not more than 3.
2. Push 1. Size is 2.
3. Push 9. Size is 3.
4. Push 3. Size is 4, which is more than 3, so pop. The smallest value, 1, leaves. The heap holds 3, 5, 9.
5. Push 7. Size is 4. Pop the smallest, 3. The heap holds 5, 7, 9.
6. Push 9. Size is 4. Pop 5. The heap holds 7, 9, 9.
7. Push 2. Pop 2 immediately, since it is the smallest. Nothing changes.
8. Push 6. Pop 6. Nothing changes.

The heap ends with 7, 9, 9, which are the three largest. Each step cost O(log k), because the heap never held more than four values.

A stream of eight values and the min-heap of three kept values after three steps

Step through the visualizer below. Each value is pushed, and the smallest value leaves whenever the heap grows past three. Watch the root at the end of each step: it is the weakest of the three values kept, which is the answer you need for the next value.

Kth largest in a stream

A stream is a sequence of values that arrives one at a time, and you must answer after each one. A search engine that keeps the best results while crawling, or a live leaderboard, works this way. The standard tool here is a class that keeps a min-heap of size k. After each value, the root is the kth largest value seen so far, as long as there are at least k values.

The class below keeps the heap. The method add pushes a value, trims the heap to size k, and returns the root, or -1 when fewer than k values have arrived. We use -1 as a signal because the values in this program are not negative. If negative values were allowed, you would return a separate flag.

Type this into the Input box. The first line is k and n, the second line holds the stream:

3 6
5 1 9 3 7 2

The output is:

-
-
1
3
5
5

The first two lines are dashes, because the stream has fewer than three values. After 9 arrives, the heap holds 1, 5, 9, and its root 1 is the third largest of 5, 1, 9. After 3 arrives, the heap holds 3, 5, 9, so the root is 3, the third largest of the four values 5, 1, 9, 3. After 7 arrives, the heap holds 5, 7, 9, and the root 5 is correct. The last value 2 is smaller than the root, so it is popped at once, and the answer stays 5.

The class stores the heap as a member. That is the reason the heap survives from one call to the next. Each call costs O(log k), so a stream of n values costs O(n log k) in total, with O(k) memory no matter how long the stream grows.

Top K frequent words

Many problems combine a counting step with a top-K step. Here the question is: which k words occur most often in a text? The counting part uses an unordered_map, which is a hash table that maps each word to a count. The top-K part uses a heap, but now the heap items are words with counts, and the order is more complicated. The strongest item is the word with the highest count. When counts are equal, the alphabetically smaller word is stronger, so the output is fixed and easy to check.

The heap must keep the k strongest words, so the root must be the weakest of them. That is the same rule as before. The comparison is written so that the top of the heap is the weakest word. The struct holds the word and its count, and the comparison function Stronger returns true when its first argument is stronger than its second.

Type this into the Input box. The first line is k and n, and the second line holds the words:

2 6
the cat the dog cat a

The output is:

cat 2
the 2

Both the and cat appear twice, and cat comes first, because it is alphabetically smaller. The words dog and a each appear once, so they are weaker than both, and the heap pops them. The result is deterministic, even though unordered_map does not promise any order when you loop over it. That is an important point: the heap's comparison decides the output, not the order the map happens to visit the words.

Notice the design of the comparison. It takes two words and decides which is stronger. A bigger count is stronger, and the tie-break makes the answer unique. If you forgot the tie-break, the heap might keep the and drop cat, depending on the order of the map, and your output would change from run to run on another machine. Always decide every tie.

Merging K sorted lists

The second classic problem is merging. Suppose you have K lists, each already sorted from smallest to largest, and you want one sorted list that contains every value. A common example is combining the results from K servers, each of which returns its own sorted results. How do you do it efficiently?

The easy method is to merge two lists at a time. Merge list 1 with list 2, then merge the result with list 3, and so on. Each merge copies its inputs, and the early results get copied again and again. With N values in total, the cost is O(N K) in the worst case. That is slow when K is large.

The heap method is faster. Keep one entry in a min-heap for each list: the current smallest unused value of that list. The root of the heap is the smallest of all the current heads. Pop it, write it to the output, and push the next value from the same list. Repeat until the heap is empty. Each value enters the heap once and leaves once, and the heap never holds more than K entries. So each of the N values costs O(log K), and the total is O(N log K).

Each heap entry needs three facts: the value, which list it came from, and its position in that list. The position lets you find the next value in the same list. A tuple holds three values together, and it compares them in order: first by value, then by list number, then by position. Comparing by value first is what makes the root the smallest current head.

Walking through an example

Take three sorted lists:

• List 0: 1, 4, 9
• List 1: 2, 8
• List 2: 0, 5, 6, 7

Start by pushing the first value of each list. The heap now holds (1, 0, 0), (2, 1, 0) and (0, 2, 0), where each entry is (value, list, position).

1. The root is (0, 2, 0). Pop it. Write 0. List 2 has a next value at position 1, which is 5. Push (5, 2, 1). The heap holds (1, 0, 0), (2, 1, 0), (5, 2, 1).
2. The root is (1, 0, 0). Pop it. Write 1. List 0 has a next value 4 at position 1. Push (4, 0, 1). The heap holds (2, 1, 0), (4, 0, 1), (5, 2, 1).
3. The root is (2, 1, 0). Pop it. Write 2. List 1 has 8 at position 1. Push (8, 1, 1).
4. The root is (4, 0, 1). Pop it. Write 4. List 0 has 9 at position 2. Push (9, 0, 2).
5. The root is (5, 2, 1). Pop it. Write 5. List 2 has 6 at position 2. Push (6, 2, 2).
6. Pop 6, write it, push 7 from list 2.
7. Pop 7, write it. List 2 has no more values, so push nothing.
8. Pop 8, write it. List 1 has no more values.
9. Pop 9, write it. List 0 has no more values.

The output is 0 1 2 4 5 6 7 8 9. Each step looked only at the root, which is the smallest of the current heads, so the output is sorted without comparing every pair of values.

Three sorted lists, their heads, and the heap that holds one head per list

The next program does exactly this. The first line of input is K. Each of the next K lines starts with a length, followed by the values of that list. A list of length 0 is allowed, and the program must handle it. Type this into the Input box:

4
3 1 4 9
2 2 8
4 0 5 6 7
0

The output is:

merged: 0 1 2 4 5 6 7 8 9

Check the empty list. Its line is 0, so lists[3] is empty, and the first loop skips it. It never enters the heap, and nothing goes wrong. If the program pushed lists[3][0] without checking, it would read past the end of an empty vector, which is undefined behaviour.

The check pos + 1 < lists[src].size() decides whether the same list has a next value. Without it, the program would push a position one past the end. This is the most common bug in merge code, and it usually shows up as garbage numbers in the output or as a crash.

What happens inside

The heap holds at most one entry for each list, so its size is at most K. The lists themselves are not copied. The program reads the values from the original vectors through indexes, and the heap holds only three numbers per entry. Memory use is therefore O(K) for the heap, plus the space of the output, which is O(N) if you store the merged result.

Each pop moves one value to the output. Each push adds the next value from the same list. The loop repeats N times in total, and each step does one pop and at most one push, each costing O(log K). The arithmetic of the heap is the only real work.

The tuple comparison is also worth a look. When two entries have the same value, the comparison goes to the list index, and then to the position. That gives a fixed order, which keeps the output stable even when values repeat. The output would be the same values in the same order even if the tie-break were different, because equal values print the same, but a fixed tie-break keeps the program predictable.

Complexity: the costs side by side

In the frequent-words row, d is the number of distinct words. In the merge row, N is the total number of values. The table shows why the heap method is the standard answer: the cost depends on log k or log K, not on the full size of the input.

Tracing the frequent-words heap

The frequent-words program first counts every word, and then it feeds each distinct word into a heap that holds at most k items. The order in which the map gives the words is not promised, so trace one possible order to see the heap work. Use k = 2 and the words the (count 2), cat (count 2), dog (count 1) and a (count 1).

1. Push the (2). The heap holds the. Its size is 1, so nothing leaves.
2. Push cat (2). The heap holds the and cat. Its size is 2, so nothing leaves. The weakest of the two is the, because cat sorts first among equal counts.
3. Push dog (1). The heap holds three words, which is more than k. The weakest of the three is dog, because its count is lowest. Pop it. The heap holds the and cat.
4. Push a (1). The heap holds three words again. The weakest is a, because its count is also 1 and it is the only word with count 1 left. Pop it. The heap holds the and cat.

The final heap holds the two strongest words. Reading the heap from strongest to weakest gives cat first, then the, which is the output. A different map order gives the same final heap, because each step keeps exactly the strongest k words seen so far. That is why the answer does not depend on the iteration order, even though the steps do.

Notice that the weakest word is always the one that leaves. At step 3 the new word dog leaves at once, because it is weaker than both words in the heap. A word that is stronger than the root would replace it, but the root here is a count-2 word. This rule is the same as the number-based top-K, with the comparison written for two fields.

Common mistakes and the messages they produce

Mistake 1: using a max-heap for the top K

A max-heap of all n values does work, but it stores everything, and it is slower for large n. The bigger error is using a max-heap of size k to keep the top k largest. Then the root is the largest value, and the program throws away the wrong values. The rule is the opposite: a min-heap of size k for the largest values.

Mistake 2: forgetting to trim the heap

If the code pushes a value and never pops when the size exceeds k, the heap grows to n. The output may still look right for small inputs, but the memory and time are no longer O(k), and large inputs run slowly. Always trim right after each push.

Mistake 3: a comparison with no tie-break

When two words have the same count, a comparison that says neither is stronger makes the result depend on the order of the input or the map. The output can change between runs on another machine. Always write a second rule for ties, and make sure it is strict.

Mistake 4: reading past the end of a list

Pushing lists[src][pos + 1] without checking the size reads one past the end. For a vector this is undefined behaviour. The symptom is often a random large number in the output. If you use at instead of [], the error is clearer:

terminate called after throwing an instance of 'std::out_of_range'
  what():  vector::_M_range_check: __n (which is 3) >= this->size() (which is 3)

The message names the index and the size. Read the two numbers. If they are equal, the index is one past the end, which is the bug.

Mistake 5: pushing an empty list's first value

An empty list has no first value. Calling lists[i][0] on it reads nothing, and the result is undefined. The fix is the if (!lists[i].empty()) check in the program above.

Mistake 6: the wrong direction in the tuple comparison

Using less instead of greater for the merge heap gives the largest head first, so the output is descending, or it is not sorted at all. Read the output carefully. A descending result for a merge means the comparator is reversed.

Debugging walkthrough: the top K is wrong

Suppose the top-K program prints the wrong values. Work through these steps in order.

1. Test with k = 1. The answer must be the maximum. If it is not, the comparison is reversed. Check that the heap is a min-heap, using greater<int>.
2. Test with k equal to n. The answer must be every value, sorted. If some values are missing, the trim is removing too many, so check the > in if (heap.size() > k).
3. Print the heap after each push. Watch the size. It must never be larger than k after the trim.
4. Compare the output with a plain sort of the same input. Sort the input, take the last k, and compare. Any difference shows you the first wrong step.

Testing with very small inputs by hand is faster than reading the code, and it finds the mistake in most cases.

Edge cases you should test

• k larger than n. The heap never reaches size k, so the kth-largest answer stays -1, and the top-K answer is simply all values in sorted order. This is correct.
• k equal to 1. The heap holds one value, so the top-K answer is the maximum, and the stream answer is the maximum so far.
• Repeated values, such as 9 9 9. The heap keeps copies, so the top-K answer may contain repeated values, and that is correct.
• All lists empty in the merge. The heap is empty from the start, the loop does nothing, and the output is merged: with nothing after it. This is correct, although a real program might print a newline only.
• Only one list. The heap holds one entry, and the merge just prints that list in order.
• Lists of very different lengths. The heap stays small, and the shorter lists leave the heap first, which is handled by the check on the next position.

Choosing a method for top K

Four methods are common. Sorting everything and reading the end is simple and costs O(n log n). It is fine when the data fits in memory and you need the full sorted order anyway. A heap of size k costs O(n log k), takes only a few lines, and works on a stream, so you never need all the data at once. Quickselect is a partition method that finds the kth value in O(n) on average. It needs all the data at once, its worst case is O(n squared), and the code is harder to get right. Counting works when the values are small whole numbers, such as scores from 0 to 100. One pass fills a count array, and reading the top values costs time proportional to the range of the values.

Each method has a price. The sort uses the most time when k is small, because it orders values you will never print. The heap is the best general answer when k is small compared with n, because it reads each value once and keeps only k of them. Quickselect is fastest on average, but a bad input order can make it slow, and it cannot answer a stream. Counting is the fastest of all when the values are small, but it breaks down when the range is huge, such as arbitrary 64-bit numbers.

In an interview, the heap is usually the answer you give first, because it is correct for streams and its cost is easy to prove. Mention quickselect when the interviewer asks for an average O(n) bound, and mention counting when the values are known to be small. Say which method you chose and why. That reasoning is what the interviewer is listening for.

A worked example with ties

Ties are where top-K answers go wrong. Take the values 4, 7, 4, 9, 7, 7 with k = 3. Sorted from largest, they are 9, 7, 7, 7, 4, 4, so the top three are 9, 7, 7. The third place is a tie between 7 and 7, but both are 7, so the answer is clear. Now change the problem so each value is a record with a name and a score, such as (ann, 7), (bob, 7), (cy, 7). With k equal to 2, the names kept depend entirely on the tie-break you choose.

With the min-heap of size k, each record is pushed, and when the heap is over size the weakest record is popped. Suppose the tie-break says the record with the smaller name is stronger. Then ann beats bob, which beats cy, when all scores are 7. With k equal to 2, the heap ends with ann and bob, whatever order the three records arrive in, so the output does not depend on the input order. That is the reason the frequent-words program compares the word as well as the count.

Write the tie-break into the comparison as a second rule, and test it with three equal values. If the output changes when you shuffle the input, the tie-break is missing.

Real-world uses of top K and K-way merges

A search engine keeps the top results while it scans each shard of its index, then merges the shard results. A log analyser keeps the top error messages by count over a day of logs. A video site keeps the trending videos for a live window of time. A database uses a K-way merge when it sorts a table too large for memory: it sorts chunks that fit, writes each sorted chunk to disk, and then merges the chunks with a heap of their current heads. That last use is why merging K lists matters outside interviews. It is how large sorts run on real machines.

Interview questions you should be ready for

How do you find the k largest elements of a large array? Use a min-heap of size k. Push each value, and pop when the size exceeds k. The cost is O(n log k) and the memory is O(k).

Why a min-heap and not a max-heap for the k largest? The root of the min-heap is the weakest value you keep, so you can compare each new value with it and drop the weakest value in O(log k).

How do you find the kth smallest element? Use a max-heap of size k, by the mirror argument.

How do you merge two sorted arrays? Use two pointers, one per array, and take the smaller head each time. That is O(n + m).

How do you merge K sorted lists? Use a min-heap of the current heads, one per list. Pop the smallest, push the next value from the same list. The cost is O(N log K).

Can you find the kth largest faster than O(n log k)? Yes. Quickselect finds it in O(n) on average, but it is harder to write correctly and its worst case is O(n squared). A heap gives a guaranteed bound.

Practice

1. Change the top-K program so it prints the k smallest values in ascending order. Which heap do you use, and which comparison changes?

2. Run the kth-largest stream program with k = 1 on 4 2 8 8 1. What does it print?

3. For the stream program, what happens when k is 3 and the first value is -5? Explain why -1 can be a bad sentinel for this program.

4. Add a k that is larger than n to the frequent-words input, such as 10 6. What does the program print?

5. Change the frequent-words comparison so that ties go to the alphabetically larger word. Predict the output for the sample input before you run it.

6. Write a program that merges two sorted arrays using two pointers and no heap. Explain why the heap version is not needed for K = 2.

7. Make the merge program print the source list number with each value, such as 0:0 1:1. Which field of the heap entry supplies the list number?

8. In the merge program, what would go wrong if you pushed the next value before printing the current one? Does the output change?

9. Give an input for the top-K program where k equals n, and explain what the output is.

10. Explain why the merge program's heap never grows beyond K entries, even if the lists are very long.

11. Why does the merge program convert src and pos to size_t or int in the size check? Write the check in a way that avoids a signed-unsigned warning.

12. A coworker says: "Top K frequent words can skip the heap and sort the counts." Give one reason the heap is better when k is small and the number of distinct words is large.

What you now know

• To find the k largest values, keep a min-heap of size k and pop whenever it grows past k. The root is the weakest value you keep.
• To find the k smallest values, use a max-heap of size k.
• A class with a heap member can answer the kth-largest query after each value in a stream, at O(log k) per value.
• Top-K with a count uses a hash map to count, then a size-limited heap with a full comparison, including a tie-break.
• Merging K sorted lists uses a min-heap of the current heads. Each pop is followed by a push from the same list if it has a next value. The cost is O(N log K).
• Check empty lists and the end of each list before you read the next value.