Kth largest: quickselect versus a heap

Why this matters

Many questions ask for a value by its rank, not by its position in a list. "Who is the third fastest runner?" "What is the fifth best selling product?" "What response time is slower than 90 percent of requests?" Each question asks for the kth largest value, where k is a number you choose.

The obvious answer is to sort the list and read the item at the right place. That always works. But sorting does more work than the question needs. You want one value, not the whole order.

This lesson shows two faster ways. The first uses a heap of size k, which is a structure from earlier in this module that keeps the largest values under control. The second is quickselect, which uses a splitting step called partitioning. You will write both, measure what they cost, and learn when each one is the right choice.

The next lesson in this module, on finding the k closest points, uses the same size-k heap idea on a different problem.

The problem with a small example

Here is a list of six numbers: 3, 2, 1, 5, 6, 4. The second largest value (k = 2) is 5. If you write the list from largest to smallest, you get 6, 5, 4, 3, 2, 1, and the second item is 5.

A value can appear more than once, and each copy counts. In the list 3, 2, 3, 1, 2, 4, 5, 5, 6, the sorted order from largest to smallest is 6, 5, 5, 4, 3, 3, 2, 2, 1. The fourth largest (k = 4) is 4.

C++ counts array positions from 0. So in a list sorted from largest to smallest, the kth largest value sits at index k - 1. The second largest is at index 1, and the first largest is at index 0.

How we measure cost

We will compare the methods by two things. The first is time, meaning how many steps the method takes. The second is memory, meaning how many extra values the method keeps while it works.

To describe time without counting every step, programmers use Big-O notation. It says how the number of steps grows when the list gets longer. Let n stand for the number of values in the list.

• O(n) means the steps grow in step with n. Ten times more values means about ten times more steps.
• O(n log n) means a bit more than O(n), because each value costs about log n steps. For a list of one million values, log n is about 20 (using base 2).
• O(n squared) means the steps grow with the square of n. Ten times more values means about one hundred times more steps.

Method one: sort the list

The simplest method sorts the list from largest to smallest and reads the value at index k - 1.

#include <algorithm>
#include <functional>
#include <vector>
 
int kthLargestBySort(std::vector<int> nums, int k) {
    std::sort(nums.begin(), nums.end(), std::greater<int>());
    return nums[k - 1];
}

The function takes nums by value. That means it works on its own copy, so the caller's list is not changed. The std::greater<int>() part tells std::sort to put larger values first.

With the list 3, 2, 1, 5, 6, 4 and k = 2, the sorted list is 6, 5, 4, 3, 2, 1. Index 1 holds 5.

The cost is O(n log n), because sorting is the expensive part. The extra memory is the copy, which holds n values.

This method is short and correct for any k. Its weakness is that it sorts all n values even though you only need one of them.

Method two: a heap of size k

The heap method asks a smaller question. It keeps only the k largest values seen so far. It uses a min-heap, which is a heap where the smallest value is always on top, for this purpose.

Why does this work? Suppose you have already read every value and kept the k largest. The smallest of those k values is the kth largest of the whole list. Every other value you kept is larger, and every value you threw away is smaller than that smallest one. So the top of the heap is the answer.

While you read the values, the heap may briefly hold k + 1 values. When it does, you remove the smallest, which can never be one of the k largest. Here is a walk through the list 3, 2, 1, 5, 6, 4 with k = 2.

At the end the heap holds 5 and 6. The top is 5, which is the second largest value. Here is the full program. It uses std::priority_queue, the standard library's heap, with std::greater<int> so that it acts as a min-heap.

#include <algorithm>
#include <functional>
#include <iostream>
#include <queue>
#include <vector>
 
// Returns the kth largest value. k = 1 means the largest.
int kthLargestBySort(std::vector<int> nums, int k) {
    std::sort(nums.begin(), nums.end(), std::greater<int>());
    return nums[k - 1];
}
 
// Keeps the k largest values seen so far in a min-heap.
int kthLargestByHeap(const std::vector<int>& nums, int k) {
    std::priority_queue<int, std::vector<int>, std::greater<int>> heap;
    for (int x : nums) {
        heap.push(x);
        if (static_cast<int>(heap.size()) > k) {
            heap.pop();  // remove the smallest value
        }
    }
    return heap.top();
}
 
int main() {
    std::vector<int> nums{3, 2, 1, 5, 6, 4};
    std::cout << "sort, k=2: " << kthLargestBySort(nums, 2) << '\n';
    std::cout << "heap, k=2: " << kthLargestByHeap(nums, 2) << '\n';
 
    std::vector<int> nums2{3, 2, 3, 1, 2, 4, 5, 5, 6};
    std::cout << "sort, k=4: " << kthLargestBySort(nums2, 4) << '\n';
    std::cout << "heap, k=4: " << kthLargestByHeap(nums2, 4) << '\n';
    return 0;
}

Compile and run it with:

g++ -std=c++20 -Wall -o kth_sort_heap kth_sort_heap.cpp && ./kth_sort_heap

The output is:

sort, k=2: 5
heap, k=2: 5
sort, k=4: 4
heap, k=4: 4

The heap method takes one push for each value. When the heap is too big, it also takes one pop. Each push and pop costs O(log k) steps, because the heap only holds about k values. So the whole method costs O(n log k). The extra memory is about k values.

When k is small, this is much cheaper than sorting. Take n = 1,000,000 and k = 3. Each heap step costs about log 3, which is under 2 steps. Sorting costs about 20 steps per value. The heap method does less work per value, and it never stores more than about 4 values at once.

Method three: quickselect

Quickselect is a method that finds the kth value without sorting everything. It is built on partitioning, the step that quicksort uses.

The partition step

Partitioning needs a pivot, which is one value chosen from the list. The partition step moves every value bigger than the pivot to the left side and every smaller value to the right side. When it is done, the pivot sits exactly where it would be in a list sorted from largest to smallest. Its position in the list tells you its rank.

Here is a short example. Use the list 3, 2, 1, 5, 6, 4 and let the pivot be the first value, 3. Each step below is what the program does.

1. Move the pivot to the end of the list. The list becomes 4, 2, 1, 5, 6, 3.
2. Scan the other values from left to right. Each value bigger than 3 moves to the front, next to the earlier ones. The values 4, 5 and 6 are bigger.
3. Put the pivot in its final place, which is index 3, and the list becomes 4, 5, 6, 3, 1, 2. Three values are before the pivot, and they are all bigger than 3.

We are looking for index 1, which is the second largest. Index 1 is less than index 3, so the answer must be in the left part, 4, 5, 6. We can ignore the values on the right. That is the key idea. Quickselect keeps only the side that holds the answer, and it repeats the partition on that side.

Continue the example. The left part is 4, 5, 6, and the first value 4 becomes the pivot. After the partition the list is 6, 5, 4, and the pivot sits at index 2. The target index 1 is less than 2, so we keep 6, 5. The pivot is the first value, 6. After the partition the list is 6, 5, and the pivot sits at index 0. The target index 1 is bigger than 0, so we keep the right part, which holds only 5. A single value is the answer. The second largest is 5.

The program

Here is the full quickselect program. It picks the pivot at a random position, which is explained after the code.

#include <algorithm>
#include <cstddef>
#include <iostream>
#include <random>
#include <utility>
#include <vector>
 
long long comparisons = 0;
 
// Partitions a[lo..hi] around the value at pivotIndex.
// Values greater than the pivot end up on the left, smaller ones on the right.
// Returns the final position of the pivot.
std::size_t partitionDesc(std::vector<int>& a, std::size_t lo, std::size_t hi, std::size_t pivotIndex) {
    std::swap(a[pivotIndex], a[hi]);
    const int pivot = a[hi];
    std::size_t store = lo;
    for (std::size_t i = lo; i < hi; ++i) {
        ++comparisons;
        if (a[i] > pivot) {
            std::swap(a[i], a[store]);
            ++store;
        }
    }
    std::swap(a[store], a[hi]);
    return store;
}
 
enum class PivotRule { First, Random };
 
// Returns the kth largest value (k = 1 is the largest).
// Requires 1 <= k <= a.size().
int quickselect(std::vector<int> a, std::size_t k, PivotRule rule, std::mt19937& rng) {
    std::size_t target = k - 1;  // position in descending order
    std::size_t lo = 0;
    std::size_t hi = a.size() - 1;
    while (true) {
        if (lo == hi) {
            return a[lo];
        }
        std::size_t pivotIndex = lo;
        if (rule == PivotRule::Random) {
            std::uniform_int_distribution<std::size_t> pick(lo, hi);
            pivotIndex = pick(rng);
        }
        std::size_t mid = partitionDesc(a, lo, hi, pivotIndex);
        if (mid == target) {
            return a[mid];
        }
        if (target < mid) {
            hi = mid - 1;
        } else {
            lo = mid + 1;
        }
    }
}
 
int main() {
    std::mt19937 rng(12345);
 
    std::vector<int> sorted(2000);
    for (std::size_t i = 0; i < sorted.size(); ++i) {
        sorted[i] = static_cast<int>(i + 1);
    }
 
    comparisons = 0;
    int a1 = quickselect(sorted, 1, PivotRule::First, rng);
    std::cout << "sorted input, first pivot: answer " << a1 << ", comparisons " << comparisons << '\n';
 
    comparisons = 0;
    int a2 = quickselect(sorted, 1, PivotRule::Random, rng);
    std::cout << "sorted input, random pivot: answer " << a2 << ", comparisons " << comparisons << '\n';
 
    std::vector<int> shuffled = sorted;
    std::shuffle(shuffled.begin(), shuffled.end(), rng);
 
    comparisons = 0;
    int a3 = quickselect(shuffled, 1, PivotRule::First, rng);
    std::cout << "shuffled input, first pivot: answer " << a3 << ", comparisons " << comparisons << '\n';
 
    comparisons = 0;
    int a4 = quickselect(shuffled, 1000, PivotRule::Random, rng);
    std::cout << "shuffled input, random pivot, k=1000: answer " << a4 << ", comparisons " << comparisons << '\n';
    return 0;
}

The comparisons counter is there only so we can measure the work. It counts how many times the partition compares a value with the pivot. A real program would not need it.

The output of this program is:

sorted input, first pivot: answer 2000, comparisons 3997
sorted input, random pivot: answer 2000, comparisons 2164
shuffled input, first pivot: answer 2000, comparisons 2168
shuffled input, random pivot, k=1000: answer 1001, comparisons 5508

The answers are correct. The counts tell a different story, which the next section explains.

Why the pivot choice matters

Each partition compares the pivot with every other value in the current range. So the first partition costs about n comparisons. The next partition works on about half the values on average, and so on. On average the total is about a small multiple of n, which is O(n).

That average needs a reasonable pivot. Here is a test on the sorted list 1, 2, ..., 2000. Each row asks for a different rank. Each row starts a new random generator with seed 1. Random pivot counts depend on the C++ library, so your numbers may differ.

sorted n=2000 k=1 first: 2000 (3997 comparisons), random: 2000 (3169 comparisons)
sorted n=2000 k=1000 first: 1001 (1999000 comparisons), random: 1001 (6229 comparisons)
sorted n=2000 k=1999 first: 2 (5994 comparisons), random: 2 (3496 comparisons)
sorted n=2000 k=2000 first: 1 (1999 comparisons), random: 1 (3494 comparisons)

Look at the row for k = 1000. With the first value as the pivot, the program makes 1,999,000 comparisons. With a random pivot it makes 6,229.

Here is why. In a sorted list the first value is always the smallest or the largest of its range. So the pivot lands at one end, and it removes only one or two values from the range on each step. Each step still compares against about as many values as the range holds. The total comes out close to n squared divided by two. For n = 2000 that is 2000 times 1999 divided by 2, which is exactly 1,999,000.

A random pivot avoids this. A sorted or nearly sorted list is no longer a trap, because a random position is unlikely to be an extreme value every time. This is why the worst case of quickselect is O(n squared), even though the average is O(n). The worst case needs a bad pivot every single time, and a random pivot makes that very unlikely.

What quickselect costs

The average time is O(n). The worst time is O(n squared). The extra memory is small, because the partition works inside the list and the loop does not use recursion. Our program copies the list once when it receives it by value, so that copy costs n values. If you pass the list by reference and accept that the caller's order will change, the copy goes away.

The C++ standard library has std::nth_element, which places the kth value in its right position with linear average complexity. In real code you should usually call that function instead of writing your own. The code in this lesson shows how the idea works.

Comparing the three ways

The table below puts the three methods side by side. Average time is what you usually see. Worst time is the slowest the method can be on any input.

Three points follow from the table.

The heap method has the same time in the average and worst cases, because every value always costs one push and sometimes one pop. It never has a bad day. Its cost depends only on k, not on how the data is ordered.

Quickselect is usually the fastest of the three when the whole list fits in memory. It rarely hits its worst case when the pivot is random.

Sorting is the most predictable and the easiest to read. Its only drawback is that it does more work than you need.

When each method wins

Sorting wins when you need the whole sorted list anyway. For example, if the program will also print the top 10 and the bottom 10, a single sort gives you both. It also wins when k is close to n, because then almost every value is needed and the extra work is small.

The heap wins when the values arrive one at a time and you cannot store them all. Here is a program that reads one million values, one per step, and keeps only the three largest. It generates the values with a simple formula, so the output is the same every time you run it.

#include <cstddef>
#include <cstdint>
#include <functional>
#include <iostream>
#include <queue>
#include <vector>
 
int main() {
    // A simple number generator that produces one value at a time.
    // The values are never stored in a list.
    std::uint32_t state = 2026;
    auto nextValue = [&state]() {
        state = state * 1664525u + 1013904223u;
        return static_cast<int>(state % 1000000u);
    };
 
    const std::size_t total = 1000000;
    const int k = 3;
    std::priority_queue<int, std::vector<int>, std::greater<int>> heap;
    std::size_t largestHeapSize = 0;
 
    for (std::size_t i = 0; i < total; ++i) {
        int x = nextValue();
        heap.push(x);
        if (static_cast<int>(heap.size()) > k) {
            heap.pop();
        }
        if (heap.size() > largestHeapSize) {
            largestHeapSize = heap.size();
        }
    }
 
    std::cout << "values read: " << total << '\n';
    std::cout << "3rd largest: " << heap.top() << '\n';
    std::cout << "most values held at once: " << largestHeapSize << '\n';
    return 0;
}

The output is:

values read: 1000000
3rd largest: 999994
most values held at once: 3

The program read a million values but never held more than three of them. Quickselect cannot do this, because its partition needs to move values around inside a list, and that list must be in memory first.

Quickselect wins when the whole list fits in memory, you want the fastest typical answer, and you can accept that the list gets reordered. A random pivot keeps the typical cost near O(n).

Common mistakes

Using a max-heap in the heap method

The heap method needs a min-heap. If you write std::priority_queue<int> heap; instead, the top is the largest value in the heap. When the heap is too big, the pop removes the largest value, so the heap keeps the k smallest values. Its top is then the kth smallest value, not the kth largest.

With the list 3, 2, 1, 5, 6, 4 and k = 2, the two smallest values are 1 and 2. The top would be 2, which is wrong. The program does not give an error message. It prints a wrong answer. Check the heap direction every time you write this method, and test one small list by hand.

Off by one with the index

The answer is at index k - 1, not k. If you write nums[k], you get the value one rank too far. With the sorted list 6, 5, 4, 3, 2, 1 and k = 2, nums[k] is 4, but the second largest is 5.

If k is larger than the list, the index goes past the end of the list. That is undefined behaviour, which means the program can crash, print garbage or appear to work. The code must check that 1 <= k <= n before it reads the list.

A warning about comparing sizes

The heap's size() function returns an unsigned number, which is a number that cannot be negative. The variable k is a signed int. Comparing them with > gives this warning when you compile with -Wall. The file below is a short example that causes it.

#include <cstddef>
#include <iostream>
#include <queue>
#include <vector>
 
int main() {
    int k = 2;
    std::priority_queue<int> heap;
    for (int x : {5, 1, 9}) {
        heap.push(x);
        if (heap.size() > k) {
            heap.pop();
        }
    }
    std::cout << heap.top() << '\n';
    return 0;
}

Compiling it with g++ -std=c++20 -Wall -o warn_sign warn_sign.cpp prints this warning:

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

The warning does not stop the program, but it points at a real risk. The fix in the lesson code is static_cast<int>(heap.size()), which turns the unsigned number into a signed one before the comparison.

Forgetting to include the heap header

If you use std::priority_queue without #include <queue>, the compiler stops with an error. Compiling this program fails:

#include <iostream>
#include <vector>
 
int main() {
    std::priority_queue<int> heap;
    heap.push(4);
    std::cout << heap.top() << '\n';
    return 0;
}

The first error is:

err_noqueue.cpp:5:10: error: ‘priority_queue’ is not a member of ‘std’

The compiler also suggests the fix in a note: add #include <queue>.

Many equal values slow quickselect down

The partition in this lesson moves only values that are strictly bigger than the pivot to the left. Values equal to the pivot go to the right, next to the smaller values. When a list has many copies of one value, the partition removes very few values at a time, and the work grows toward the worst case.

Here is a test with 1000 values, all equal to 7, asking for k = 500 with a random pivot. The count depends on the C++ library, so your number may differ.

1000 equal values, k=500, random pivot: answer 7, comparisons 374750

The answer is correct, but 374,750 comparisons is about 375 times n. A version that splits the list into three parts (smaller, equal, bigger) avoids this problem. Check how your library's std::nth_element behaves on data like this before you rely on it.

Changing the caller's list by accident

The quickselect function in this lesson takes the list by value, so the caller's list stays as it was. If you change the parameter to std::vector<int>& a to save the copy, the partition moves values around inside the caller's list. The function still returns the right answer, but the caller now sees a list in a different order. Decide on purpose whether the order matters to the rest of your program.

Practice

Exercise 1. Use the heap method with k = 3 on the list 4, 1, 7, 3, 8, 5. Write the heap after each value is read. What is the answer?

Show the answer

Keep only the three largest values. After 4, 1 and 7, the heap holds 1, 4, 7. Push 3, which makes four values, then pop the smallest, 1. The heap holds 3, 4, 7. Push 8, pop 3. The heap holds 4, 7, 8. Push 5, pop 4. The heap holds 5, 7, 8. The top is 5. To check, sort the list from largest to smallest: 8, 7, 5, 4, 3, 1. The third value is 5.

Exercise 2. A file holds 10 billion numbers. It is too big to load into memory, and you can read it one number at a time. You need the 100th largest number. Which method do you use, and why?

Show the answer

Use the heap of size k. It keeps 100 values in memory and reads each number once, so its memory use does not depend on the file size. Sorting needs the whole file in memory. Quickselect needs the whole list so that it can partition it.

Exercise 3. The heap method in this lesson returns the kth largest value. Change it so that it returns the kth smallest value instead. Which line changes?

Show the answer

Remove std::greater<int> from the type, so the priority queue becomes a max-heap with the largest value on top. Keep the pop() when the size is larger than k, and keep return heap.top();. The heap now keeps the k smallest values, and its top is the largest of those, which is the kth smallest value overall.

Exercise 4. A list has 2000 values in sorted order from smallest to largest, and you ask for k = 1000. Explain why quickselect with the first value as the pivot makes about 2 million comparisons, and why a random pivot makes far fewer.

Show the answer

With the first value as the pivot, the pivot is always the smallest or largest value of the current range. It lands at one end and removes only one or two values from the range. The range shrinks slowly, and each step compares the pivot with about as many values as the range holds. The total is close to n squared divided by two, which is 1,999,000 for n = 2000. A random pivot is very unlikely to be an extreme value every time, so the ranges shrink quickly and the total stays near a small multiple of n. The measured count was 6,229.

Exercise 5. A program needs the 3rd largest of a list of 1,000,000 values that is already in memory. It only needs that one value. Which method uses less extra memory: the heap method or quickselect? Give the reason in one sentence.

Show the answer

The heap method uses less extra memory, about 4 values. Quickselect works on that list in place, so it adds only a few variables if you pass the list by reference. Quickselect is usually the faster choice here, if you accept that the list gets reordered.

Exercise 6. Write the first partition step by hand for the list 5, 9, 2, 7 with the first value as the pivot. Which values end up on the left of the pivot, and at what index does the pivot finish?

Show the answer

Swap the pivot 5 with the last value, which gives 7, 9, 2, 5. Scan the other values. 7 and 9 are bigger than 5, so they stay at the front. The value 2 is smaller. Swap the pivot into place, which gives 7, 9, 5, 2. The pivot 5 finishes at index 2. The values on its left, 7 and 9, are both bigger than 5. The value 2 is on its right.

Summary

• Sorting finds the kth largest value in O(n log n) time, and it is the simplest method to write and read.
• A min-heap of size k keeps only the k largest values seen so far. Its top is the answer, and it costs O(n log k) time with about k values of memory.
• Quickselect uses partitioning to keep only the side of the list that holds the answer. Its average time is O(n), but its worst case is O(n squared), so the pivot should be random.
• The heap method is the right choice when values arrive as a stream that cannot be stored, or when k is small compared with a very large list.
• Quickselect is usually the fastest choice when the whole list fits in memory, but it reorders the list, and many equal values slow it down.
• In real C++ code, std::nth_element gives quickselect's average speed without you writing the partition yourself.