Kth largest in Python: quickselect versus heapq

Why this matters

A common interview question is this: given a list of numbers, find the kth largest. The letter k is a number you pick, such as 1, 2 or 3. When k is 1, you want the largest number. When k is 2, you want the second largest, and so on.

Most people first answer by sorting the list and reading the right position. That answer is correct. The interviewer then asks, "Can you do it faster, or with less memory?" That follow-up is the point of this lesson. You will write three answers: sorting, a heap, and quickselect. Each one has a different cost, and each one wins in different situations.

The earlier lessons in this module showed how a heap keeps its smallest item on top. Python's heapq module does exactly that. This lesson uses that property to answer a real question. Quickselect does not use a heap. It uses the split-and-keep-one-side idea from quicksort.

The problem in plain words

A few terms first.

A list of numbers has n items. n is the number of items, and using a letter for it lets us state a cost that works for any list.

A cost here means how much work a method does as the list grows. We write costs in Big-O notation. O(n) means the work grows about in step with the list: double the list, and the work roughly doubles. O(n log n) means the work grows a little faster than that. The word log here means the number of times you can cut a number in half before you reach 1. For 1000, that is about 10 halvings.

Take the list [3, 2, 1, 5, 6, 4] and k equal to 2. Sort it in your head from smallest to largest: 1, 2, 3, 4, 5, 6. The largest is 6, which is the first largest (k = 1). The second largest is 5 (k = 2). So the answer is 5.

When values repeat, each copy counts as its own entry. In [2, 1, 2, 2, 3], sorted from largest to smallest is 3, 2, 2, 2, 1. The second largest is 2, even though 2 appears three times.

We assume k is at least 1 and at most the number of items. A k of 0 or a k bigger than the list has no answer.

Method 1: sort, then read the answer

If you sort the list from smallest to largest, the kth largest number sits at a known position. Python counts positions from 0, which is the first position. The largest number is at index n - 1. The second largest is at index n - 2. In general, the kth largest is at index n - k.

Python also lets you count from the end. Index -1 is the last item, -2 is the one before it, and so on. So index n - k is the same as index -k.

For the list above, n is 6 and k is 2, so the index is 4, or -2. The sorted list is [1, 2, 3, 4, 5, 6], and index 4 holds 5. That matches.

Here is the code. sorted() returns a new list, so the caller's list keeps its original order.

No input is needed. The output is:

k=2: 5
k=1: 6
k=6: 1
original: [3, 2, 1, 5, 6, 4]

The last line shows that the original list did not change. sorted() made a new list, and the sort ran on that copy. The list method nums.sort() sorts the list itself, which is faster when you do not need the original order.

Sorting costs O(n log n) time. The new list costs O(n) extra memory. If you do not need the original order, you can call nums.sort() and skip the copy.

The sorting method is short and hard to get wrong. It also has one more advantage. If you need the answer for many values of k on the same data, you sort once and read any position you like.

Method 2: keep the k largest in a min-heap

Now a different idea. Do not sort everything. Keep only the k largest numbers seen so far.

A min-heap is a good tool for this. It keeps its smallest item on top. Python's heapq module gives you a min-heap inside a plain list. heapq.heappush(heap, x) adds an item. heapq.heappop(heap) removes and returns the smallest item. heap[0] reads the smallest item without removing it.

Adding or removing an item costs about log k steps, because the heap never holds more than about k items. The method works like this:

1. Push the next number onto the heap.
2. If the heap now holds more than k numbers, pop the smallest with heappop.
3. Repeat for every number in the list.

At the end, the heap holds the k largest numbers. Its top, heap[0], is the smallest of those k numbers. That top is the kth largest of the whole list.

Why remove the smallest? A new number only needs to beat the smallest kept number to get in. The smallest kept number is the first one to go. A min-heap keeps that smallest number on top, so finding it is fast.

Here is a trace with k = 2 on [3, 2, 1, 5, 6, 4]. Each line shows the heap after one step. The heap prints its items in its internal list order. Only the first item, heap[0], is guaranteed to be the smallest.

add 3 -> [3]
add 2 -> [2, 3]
add 1 -> [1, 3, 2]  pop removes 1 -> [2, 3]
add 5 -> [2, 3, 5]  pop removes 2 -> [3, 5]
add 6 -> [3, 5, 6]  pop removes 3 -> [5, 6]
add 4 -> [4, 6, 5]  pop removes 4 -> [5, 6]
answer = heap[0] = 5

Read the third line closely. After adding 1, the heap holds three numbers. That is more than k = 2, so heappop removes the smallest, which is 1. The heap then holds 2 and 3. The rest of the trace follows the same rule. At the end the heap holds 5 and 6, and its top is 5.

Here is the code for this method.

The output is:

k=2: 5
k=1: 6
k=6: 1
k=3 on 6 numbers: 10
heapq.nlargest: [6, 5]

The fourth line uses a new list, [7, 10, 4, 3, 20, 15]. Sorted, it is 3, 4, 7, 10, 15, 20. The third largest is 10, which matches.

The last line uses heapq.nlargest, a function in the standard library. It returns the k largest items, biggest first. It does the same job as the loop above in one call. Writing the loop yourself is useful practice, and the call is handy when you only need the answer.

The heap never holds more than k + 1 numbers at once. So each push or pop costs about log k steps. For the whole list the cost is O(n log k) time and O(k) memory.

To see what that means in practice, take a list of one million numbers and k = 3. Sorting costs about log n, which is about 20 steps per number. The heap costs about log k, which is about 1.6 steps per number. In the big-O count, the heap does far less work when k is small.

Method 3: quickselect

Quicksort is a common sorting method. It picks one number, called the pivot. It splits the list into numbers smaller than the pivot and numbers larger than the pivot. Then it sorts each side.

Quickselect uses the same split. It does not sort. After each split, it keeps only the side that holds the answer, and it throws the other side away. Less work is done each round.

Here is the idea with [3, 2, 1, 5, 6, 4], where we want the second largest. That is sorted position 4, counting from 0 at the smallest.

Round 1. Pick 3 as the pivot. Split the list into three groups. The smaller group is [2, 1]. The equal group is [3]. The larger group is [5, 6, 4]. Once sorted, the smaller group takes positions 0 and 1, the equal group takes position 2, and the larger group takes positions 3, 4 and 5. Position 4 is in the larger group. We throw away the other two groups.

Round 2. The larger group is [5, 6, 4], and it still holds the target. Its first item is at position 3, so the target is the second item in this group. Pick 5 as the pivot. The smaller group is [4] and sits at position 3. The equal group is [5] and sits at position 4. The target is position 4, so the answer is 5.

The first round looked at 6 numbers, and the second looked at 3. The list gets smaller each round, and we never sort the whole list.

Why three groups and not two? When values repeat, a two-way split can put every copy on one side. The list then shrinks by only one value each round. A three-way split puts all copies equal to the pivot into their own group, so they are finished at once.

Why a random pivot? If the pivot is always the largest or smallest number in the range, each round removes only one number. Picking the pivot at random makes that bad pattern very unlikely, even when the input is already sorted.

random.Random is Python's random number generator. A seed such as random.Random(7) makes it give the same numbers on every run, so the output below stays fixed. rnd.randint(lo, hi) returns a whole number from lo to hi, and both ends can be picked. Java's nextInt(n), which is the nearest match, stops at n minus 1, so check the ends when you move code between the two.

Here is the code. The select function works on a list in place. It returns the value that would sit at a given sorted position. kth_largest calls it with the position n - k.

The output is:

k=2: 5
k=1: 6
k=6: 1
k=2 with duplicates: 2
array after select: [1, 2, 3, 4, 5, 6]

The last line matters. select reorders the list it is given. The list is no longer in its original order, and it is not guaranteed to be sorted. This run happens to end sorted. Only the item at the target position is guaranteed to be right. kth_largest avoids the problem by copying the list first.

The while loop has one job. It keeps lo and hi as the range where the answer can be. Each round the range shrinks, either to the smaller group or to the larger group. When the range has one item, that item is the answer.

Quickselect costs O(n) time on average. Each round looks at every number in the current range once. With a random pivot, the range shrinks by about half on average, so the total work adds up to a fixed multiple of n. In the worst case, every pivot is the smallest or largest value, the range shrinks by one each round, and the work is O(n^2). Random pivots make that case very rare.

Quickselect needs no extra list. It uses a few variables and works inside the list it was given. The cost is that the list gets reordered.

Measuring the difference

Big-O is a rough count. To see the real behaviour, the next program counts steps. A step is one time the split loop looks at one number. Counting steps gives the same result on every computer with the same Python version. A timing in seconds would change from machine to machine.

The program uses a list of 2000 numbers and asks for the middle value, so k is 1000 in the sorted order. It runs quickselect in four cases. Each case also checks its answer against sorting. The global line lets the functions update the step counter that belongs to the whole program.

The output is:

sorted input, last element as pivot: steps=1500500 answer=1000 correct=True
sorted input, random pivot: steps=7710 answer=1000 correct=True
random input, random pivot: steps=6009 answer=500935 correct=True
all equal, random pivot: steps=2000 answer=5 correct=True
random tests agreeing with sort: 1000 of 1000

Read the rows one at a time.

The first row uses sorted input with the last element as the pivot. On sorted input, the last element is always the largest. Each round then removes only one number from the range. The count of 1,500,500 is below n times n divided by 2, which is 2,000,000. The search ends once the pivot lands on the target, so the rounds run from size 2000 down to size 1001.

The second row uses the same sorted input with a random pivot. The count drops to 7,710, about 3.9 times n. The random pivot removed the bad pattern.

The third row uses random input with a random pivot. The count is 6,009, about 3 times n. This is the typical case.

The fourth row uses a list where every number is 5. The three-way split puts all of them into the equal group in one round. The count is exactly n.

The last line runs 1000 small random lists through select and compares each answer with sorting. All 1000 agree. That check is the reason to trust the code.

The random-pivot counts depend on the random numbers that random.Random gives for each seed. A different generator, or a different Python version that changes those numbers, gives different counts. The pattern stays the same: the random pivot costs a few times n, and the last-element pivot on sorted input costs about n squared divided by 2.

Comparing the three methods

The table shows the costs side by side. The memory column counts extra memory, beyond the input list.

These rules help you choose.

If the numbers arrive one at a time, and you cannot keep them all, use the heap. The heap is the only one of the three that never needs the whole list at once. Suppose a log file holds 10 billion response times, and you want the 100th largest. A heap of size 100 reads one number at a time and stores 101 numbers at most.

If k is small compared with n, the heap is fast and uses little memory. Its cost grows with log k, which stays small even for big lists.

If you have one list in memory, one answer to find, and no need for a fixed worst-case time, quickselect is usually the fastest on average. It also needs no extra list. The cost is the reordering, and the rare slow case.

If you need the answer for many values of k, sort once. Every answer after that is one read.

If the list is small, or you want the code that is least likely to have a bug, sort. For ten numbers all three methods are instant, and the sort needs the fewest lines.

If you need a guaranteed time limit, the sort and the heap both have a fixed worst case. Basic quickselect does not.

Common mistakes

The examples below are short programs with real output from Python 3.14. The folder names and line numbers in the error messages depend on your file.

Off by one in the index

Python lists start at index 0 with the smallest number. A common mistake is to read ordered[k], which counts from the smallest. The code below is wrong.

def kth_largest_wrong(nums, k):
    ordered = sorted(nums)
    return ordered[k]
 
 
nums = [3, 2, 1, 5, 6, 4]
print("k=2 returns", kth_largest_wrong(nums, 2))
print("k=6 returns", kth_largest_wrong(nums, 6))

Running it gives this output:

k=2 returns 3
Traceback (most recent call last):
  File "/path/to/your_file.py", line 8, in <module>
    print("k=6 returns", kth_largest_wrong(nums, 6))
                         ~~~~~~~~~~~~~~~~~^^^^^^^^^
  File "/path/to/your_file.py", line 3, in kth_largest_wrong
    return ordered[k]
           ~~~~~~~^^^
IndexError: list index out of range

The first line is a silent wrong answer. For k = 2 the code returns 3 with no error. The second part is the error for k = 6. The message says the code asked for index 6 in a list of length 6. Valid indexes run from 0 to 5. The fix is ordered[-k], or ordered[len(ordered) - k].

A min-heap used where a max-heap is needed

The heap must keep the k largest numbers, so its top must be the smallest of them. A max-heap does the opposite. It keeps the smallest numbers, and its top is the largest of those. heapq only makes min-heaps. To build a max-heap, a common trick is to push negative numbers. The code below pushes -x but then reads the top as if it were a min-heap of the real values. It is wrong.

import heapq
 
 
def kth_largest_wrong_heap(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, -x)
        if len(heap) > k:
            heapq.heappop(heap)
    return -heap[0]
 
 
print("k=2 returns", kth_largest_wrong_heap([3, 2, 1, 5, 6, 4], 2))

The output is:

k=2 returns 2

There is no error message. The program just prints a wrong answer. The correct answer is 5. The only way to catch this kind of bug is a test with an answer you already know. Use a short list, work out the answer on paper, and compare.

Taking the result of heappop on an empty heap

heapq.heappop raises an error when the heap is empty. The next program does this on purpose.

import heapq
 
heap = []
heapq.heappush(heap, 7)
first = heapq.heappop(heap)
print("first =", first)
second = heapq.heappop(heap)
print("second =", second)

The output is:

first = 7
Traceback (most recent call last):
  File "/path/to/your_file.py", line 7, in <module>
    second = heapq.heappop(heap)
IndexError: index out of range

The message is short. The heap is empty, so there is no smallest item to return. Check if heap: before you pop.

Quickselect reorders the caller's list

The select function changes the list it receives. The snippet below needs the select function from the quickselect section. Put that function above this code, after import random.

import random
 
scores = [90, 75, 60, 88, 95]
print("before:", scores)
third = select(scores, len(scores) - 3, random.Random(3))
print("third largest:", third)
print("after:", scores)

The output is:

before: [90, 75, 60, 88, 95]
third largest: 88
after: [60, 75, 88, 90, 95]

The answer is right, but the list is now in a different order. Any later code that expects the original order will fail in a quiet way. Copy the list first, as kth_largest does, or write a comment that says the function reorders its input.

A fixed pivot on sorted input

The experiment above showed the cost. With the last element as the pivot, sorted input took 1,500,500 steps. With a random pivot, the same input took 7,710. The fix is to pick the pivot at random, as the code does with rnd.randint.

Practice

Exercise 1

Use the sorting method to find the 4th largest value in [9, 1, 7, 7, 3].

Show the answer

Sorted from smallest to largest, the list is 1, 3, 7, 7, 9. The length is 5, so the index from the end is -4, which is the same as index 1. Index 1 holds 3. The answer is 3. Check it by counting down from the largest: 9, then 7, then 7, then 3. That is the fourth.

Exercise 2

Write kth_smallest(nums, k) with a max-heap of size k. In Python, a max-heap is a min-heap of negated numbers. Test it with [3, 2, 1, 5, 6, 4] and k = 2. Explain in one sentence why a max-heap is the right choice here.

Show the answer

The max-heap keeps the k smallest numbers, and its top is the largest of them. At the end, that top is the kth smallest. The method is the same as the kth largest one, with the sign of each number flipped.

import heapq
 
 
def kth_smallest(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, -x)
        if len(heap) > k:
            heapq.heappop(heap)
    return -heap[0]
 
 
print("kth_smallest k=2:", kth_smallest([3, 2, 1, 5, 6, 4], 2))

Running it prints kth_smallest k=2: 2. The two smallest numbers are 1 and 2, and the top of the heap is 2.

Exercise 3

Quickselect is given the list [5, 5, 5, 5] and k = 2. How many steps does the three-way split take, and why?

Show the answer

It takes 4 steps, one for each number. The pivot is 5, so every number is equal to the pivot and goes into the equal group. The equal group covers the target index, so the function returns 5 after one pass. The experiment in this lesson showed the same effect with 2000 equal numbers, which took exactly 2000 steps.

Exercise 4

A file holds 10 billion numbers, too many to fit in memory. You want the 100th largest. Which method do you use, and why not the other two?

Show the answer

Use the min-heap of size 100. It reads one number at a time and keeps at most 101 numbers. Sorting needs all the numbers in memory. Quickselect also needs the full list in memory, because it reorders it.

Exercise 5

Compare the work for a list of one million numbers when k is 3. How much more work does sorting do than the heap method, in the big-O count?

Show the answer

The heap does about log2(3), which is about 1.6 steps per number. Sorting does about log2(1,000,000), which is about 20 steps per number. In the big-O count, sorting does about 12.6 times as much work per number. Big-O drops constant factors, so the real speed gap depends on the code and the machine.

Summary

• Sorting gives the answer in O(n log n) time and is the shortest code. sorted() returns a new list, so the input stays the same.
• A min-heap of size k gives the answer in O(n log k) time and O(k) memory, and it works when numbers arrive one at a time.
• Quickselect averages O(n) time with a random pivot, has a rare O(n^2) worst case, and reorders the list it is given.
• Pick the method by what you have: a stream, a memory limit, one list, or many questions on the same data.
• Check the index -k, the sign of the numbers in a max-heap, the IndexError that heappop raises on an empty heap, and whether the list order matters.