Quickselect in Python

Why this matters

Many programs need one value from a list by its place in sorted order, without caring about the rest of the list. Examples are the median price in a list of prices, the tenth highest score in a game, or the smallest distance in a set of distances.

The simple way is to sort the whole list and read the place you want. Sorting costs about n log n steps for n items. Here log n is how many times n can be cut in half before reaching 1, about 10 for 1000 items. Quickselect gets the same value in about n steps on average. It does this by throwing away the part of the list that cannot hold the answer after each round.

Quickselect uses the partition step from the quick sort lesson. That lesson explains partitioning in full. This lesson uses the same partition step and does not repeat the reasons behind it.

The problem in plain words

A rank is a position in sorted order. This lesson counts ranks from 0, so rank 0 is the smallest item. Programmers call this a zero-based index. Python lists also count positions from 0, so the rank and the list index are the same number.

Take the list [10, 7, 8, 9, 1, 5]. Sorted, it is 1, 5, 7, 8, 9, 10. The item at index 0 is 1. The item at index 2 is 7. So "the value at index 2 of the sorted list" is 7.

The input to quickselect is a list and a number k. The output is the value that would sit at index k if the list were sorted. The list itself does not need to be sorted, and quickselect does not sort it.

Watch out for the wording. People often say "third smallest" when they mean index 2. In this lesson, the third smallest value is the one at k equal to 2.

The idea

Quickselect repeats three steps on a range of the list. A range is the part from index lo to index hi, both included.

1. Partition. Pick a pivot, which is one item in the range. Move the smaller items to its left and the bigger items to its right. The pivot now sits at its final sorted position. Call that position p.
2. Compare. Look at p and k.
   • If p equals k, the pivot is the answer. Stop.
   • If k is less than p, the answer is on the left. The right side can be thrown away.
   • If k is bigger than p, the answer is on the right. The left side, with the pivot, can be thrown away.
3. Repeat. Run the same steps on the side you kept. When the range has one item, that item is the answer.

The key difference from quick sort is in step 2. Quick sort sorts both sides, so it must visit both. Quickselect keeps one side only, so each round does less work than the round before.

Here is the idea on the list [10, 7, 8, 9, 1, 5] with k equal to 2. The pivot is 5. After the partition, 5 sits at index 1. Index 2 is bigger than 1, so the left side is no longer needed. The new range is indexes 2 to 5. That range holds 8, 9, 10 and 7.

The cost of each round is about the size of the range, because the partition looks at each item once. The first round costs n. If the pivot lands near the middle, the next range is about half as big. So the costs form a series: n, then about n divided by 2, then about n divided by 4, and so on. A series like this adds up to about 2n in total. This is the reason the average cost is linear, which is covered in a later section.

The partition step

The partition used here is the Lomuto partition from the quick sort lesson. It uses the last item of the range as the pivot. It keeps a marker i that points to the last item of the "smaller" zone, and a second marker j that scans the range.

def partition(a, lo, hi):
    pivot = a[hi]
    i = lo - 1                          # end of the "smaller" zone
    for j in range(lo, hi):
        if a[j] < pivot:
            i += 1
            a[i], a[j] = a[j], a[i]
    a[i + 1], a[hi] = a[hi], a[i + 1]   # put the pivot in its place
    return i + 1

The function range(lo, hi) produces the numbers from lo up to, but not including, hi.

The line a[i], a[j] = a[j], a[i] swaps two items. Python builds the pair on the right first, then assigns it, so no temporary variable is needed. If you know Java, this replaces the swap helper method.

The function returns the index where the pivot ended up. That index is the value p from the idea section.

The loop

The loop uses two variables, lo and hi, to mark the range that still holds the answer. Each round calls partition on that range, then moves lo or hi so the range shrinks.

def quick_select(a, k):
    lo = 0
    hi = len(a) - 1
    while lo < hi:
        p = partition(a, lo, hi)
        if p == k:
            return a[p]
        if k < p:
            hi = p - 1                  # keep the left side
        else:
            lo = p + 1                  # keep the right side
    return a[lo]                        # one item left, it is the answer

The loop is written with while and not with a recursive call. A recursive version, one that calls itself, would make one call per round. On sorted input, the rounds can be as many as the number of items, so the recursion can run very deep. The loop uses a fixed amount of memory no matter how many rounds it takes.

Python does limit recursion. The limit is 1000 nested calls by default, and the function sys.getrecursionlimit() reports it. The program below runs a recursive version on sorted lists of two sizes.

The output is:

recursion limit: 1000
n=500: 0
n=2000: RecursionError: maximum recursion depth exceeded

The recursive version works for 500 items and fails for 2000, because a sorted list with the last item as pivot needs about one call per item. The while loop above has no such limit.

The loop ends in one of two ways. Either p equals k and the function returns the pivot, or the range shrinks to one item and the last line returns it. Both cases return the value at index k of the sorted list.

Worked example 1: a trace of one search

Find the value at index 2 of [10, 7, 8, 9, 1, 5]. The answer should be 7. The program below is the loop with a print statement added after each partition. That line prints the state so you can follow the rounds.

The output is:

step 1: range 0..5, pivot 5 ends at index 1, array [1, 5, 8, 9, 10, 7]
step 2: range 2..5, pivot 7 ends at index 2, array [1, 5, 7, 9, 10, 8]
index 2 in sorted order: 7

Read step 1 slowly. The partition moved 1 to the left of the pivot and left 8, 9, 10 and 7 to the right. The pivot 5 is now at index 1. Its place is final: one item is smaller and the rest are bigger. The target k is 2, which is bigger than 1, so the code sets lo to 2 and drops indexes 0 and 1.

Step 2 works on the range 2 to 5, which holds 7, 8, 9 and 10 in some order. The pivot is 7. The partition places it at index 2. That index equals k, so 7 is the answer.

Notice that the list is not sorted at the end. Only the pivots are in their final places. The other items are only sorted into two groups that are split by the pivot. That is all quickselect needs.

Worked example 2: checking against a sort

A new function needs a test. The simplest test compares its answer with the answer from sorted, which is known to be correct. The program below makes 2000 random lists with many repeated values. It checks every k it picks. It also contains a version of quick_select with a check on k. That check is covered in the mistakes section.

The output is:

passed 2000 random tests with duplicates
sorted scores: [40, 61, 72, 79, 88, 95, 95]
largest (k=0): 95
second largest (k=1): 95
third largest (k=2): 88
scores after the calls: [72, 95, 40, 88, 95, 61, 79]

The first line shows that all 2000 answers matched the sorted answer. random.Random(7) makes a random number generator with a fixed seed, so the same numbers come out on every run. The test uses small values and repeats, because repeats are where partition code often breaks. Two equal items can be placed on either side of the pivot. The Lomuto partition handles this by putting equal items on the right side. The loop still works because it only needs the value at the pivot's final place.

The last part of the output shows the third use of the function. kth_largest asks for index len(a) - 1 - k instead of index k. Counting from the top, k equal to 0 means the last index of the sorted list. So k equal to 0 gives the largest value, k equal to 1 gives the second largest, and so on.

The last line shows the list is the same as before the calls. kth_largest passes a copy, a[:], into quick_select. Without the copy, the partition would rearrange the caller's list.

Worked example 3: cost of the search

The cost of quickselect depends on the luck of the pivot. The next program counts comparisons. A comparison is one test of a[j] < pivot inside the partition. The count measures the work done.

The program runs on two kinds of input. The first is a random list. The second is an already sorted list 0, 1, 2, and so on, using the last item as the pivot. Both searches use k equal to n divided by 2, which asks for the middle value. The // operator in n // 2 divides and drops the fraction, which gives a whole number that can be used as an index.

The counter is a global variable, a variable defined at the top level of the file. Inside a function, the line global comparisons tells Python that comparisons += 1 changes the top-level variable. Without it, Python would treat the name as a new local variable and fail with an UnboundLocalError.

The output is:

| n | random input | sorted input |
|---|---|---|
| 1000 | 4459 | 374750 |
| 2000 | 5656 | 1499500 |
| 4000 | 13565 | 5999000 |
| 8000 | 39900 | 23998000 |
random pivot, sorted input, n=8000: 21994 comparisons
last-item pivot, sorted input, n=8000: 23998000 comparisons

These numbers come from the counting program above, which uses the fixed seed 42 with Python 3.14, so every run on that version prints the same numbers. Other Python versions or other seeds can print different random numbers. The random column changes with the random input, so its exact values are not important. This table is one run with one seed. In this run, going from 1000 to 8000 items, which is 8 times as many, made the random cost about 9 times bigger. One run is not enough to read a pattern from that column. The average-case section below gives the expected pattern.

The sorted column is clear. Each time n doubles, the cost grows by about 4 times. That is the pattern of n squared, which means the cost is about 3/8 of n squared for this input.

The reason is the pivot choice. The last item is the largest item in a sorted range. The partition puts the pivot at the far right, so the next range is only one item smaller. The search removes one item per round, and each round costs about as many comparisons as there are items left. The rounds cost n minus 1, then n minus 2, and so on, down to the middle. For k equal to n divided by 2, that sum is about 3/8 of n squared, which matches the table.

Average and worst case

The average case is the cost you expect for a typical input, and the worst case is the cost for the worst possible input. These are two different questions. Quickselect has a good average case and a bad worst case.

Average case. When the pivot is random, each round splits the range at a random place. Half of the time the range shrinks a lot. The expected total work is a fixed multiple of n. A standard analysis shows the expected number of comparisons is at most 4n, no matter how big n is. So the average time is linear, written O(n). That means the steps grow in step with n: double the items and the steps about double.

Worst case. The worst case happens when every pivot is the smallest or the largest item of its range. Then each round removes only one item. The rounds cost n, then n minus 1, then n minus 2, and so on. The total is about n squared divided by 2, so the worst time is O(n squared).

The worst case happens with a simple pivot rule on sorted or nearly sorted input. The sorted column in the table shows it. The fix is a random pivot. The change is one extra step before each partition: it swaps a random item in the range into the last place. Then the partition uses that item as the pivot.

r = rng.randrange(lo, hi + 1)           # random index from lo to hi
a[r], a[hi] = a[hi], a[r]               # move it to the last place
p = partition(a, lo, hi)

With a random pivot, the same sorted input takes about 22 thousand comparisons instead of about 24 million. This does not make the worst case impossible. A bad pivot is still possible by chance, but it is very unlikely over many rounds.

Common mistakes

Mistake 1: treating k as a count instead of an index. The function takes an index that counts from 0. If you pass 3 to get the third smallest, you get the fourth smallest. For [10, 7, 8, 9, 1, 5], the call with k equal to 3 returns 8. The sorted list is 1, 5, 7, 8, 9, 10, and index 3 holds 8. Check the index before you call the function.

Mistake 2: no check on k. If k is outside the list, the loop does not fail at once. It keeps moving the range until one item is left, then it returns that item. So the answer looks normal but is wrong. A small test shows it. The program below calls a version without a check with k equal to 5 on the list [3, 1, 2], which has only indexes 0 to 2. It also calls it with k equal to -1.

The output is:

no check, k=5 returns: 3
no check, k=-1 returns: 1

The program prints values with no error. Negative numbers are no protection here. In a list, a[-1] means the last item, but quickselect uses k as a rank from the bottom, so -1 is smaller than every rank. The search keeps the left side each round and returns the smallest item, which is 1.

The check in the version from worked example 2 stops this. Calling that version with k equal to 5 on the same list, with no code to catch the error, prints this:

Traceback (most recent call last):
  File ".../s14_9.py", line 28, in <module>
    print(quick_select([3, 1, 2], 5))
          ~~~~~~~~~~~~^^^^^^^^^^^^^^
  File ".../s14_9.py", line 14, in quick_select
    raise ValueError(f"k must be from 0 to {len(a) - 1}, got {k}")
ValueError: k must be from 0 to 2, got 5

The last line is the message the check writes, with the type of error, ValueError, before it. The lines starting with File make up the traceback, the list of calls that were running when the error happened. The path before the file name depends on where you saved the file, so yours will differ.

Mistake 3: an empty list. An empty list has no index 0. Without the check, the loop does not run, and the last line tries to read index 0. The program below calls the version without a check on an empty list.

print(quick_select_no_check([], 0))

Python stops with this message:

Traceback (most recent call last):
  File ".../s14_5.py", line 26, in <module>
    print(quick_select_no_check([], 0))
          ~~~~~~~~~~~~~~~~~~~~~^^^^^^^
  File ".../s14_5.py", line 23, in quick_select_no_check
    return a[lo]
           ~^^^^
IndexError: list index out of range

The check from worked example 2 catches this case too, because any k is outside an empty list. With the check, the same call stops with "k must be from 0 to -1, got 0". That message is clearer, because it names the bad value and the valid range.

Mistake 4: moving lo or hi to the pivot instead of past it. The lines hi = p - 1 and lo = p + 1 leave the pivot out of the next range. The pivot is already in its final place. If you write lo = p instead of lo = p + 1, the range can stop shrinking. When p equals lo, the code sets lo to the same value again. With repeated values, that can happen in every round. For [1, 1, 1] with k equal to 2, p is 0 each time, so the loop never ends.

The program below runs the broken loop for only five rounds, so it stops. The real loop has no such limit, so it would run forever.

The output is:

round 1: pivot ends at 0, range 0..2
round 2: pivot ends at 0, range 0..2
round 3: pivot ends at 0, range 0..2
round 4: pivot ends at 0, range 0..2
round 5: pivot ends at 0, range 0..2

The range stays at 0 to 2 in every round. To stop a program like this, press Ctrl+C in the terminal.

Mistake 5: changing the caller's list without knowing it. The partition moves items inside the list it gets. A call that looks like a read, such as quick_select(scores, 2), changes the order of scores. The program below shows it.

The output is:

answer: 7
scores after the call: [1, 4, 7, 9]

If the caller needs the original order, pass a copy with scores[:] or list(scores), as the examples in this lesson do.

Practice

Each exercise has its answer below it. Try each one on paper or in a scratch file before you read the answer.

Exercise 1: trace by hand

Find the value at index 1 of [4, 2, 9, 6, 5] by hand. Write the list after each partition, the pivot, and the new range.

Show the answer

The first range is indexes 0 to 4, and the pivot is 5. Items smaller than 5 are 4 and 2. The partition moves them left, so the list becomes [4, 2, 5, 6, 9], and the pivot 5 is at index 2. Index 1 is less than 2, so the new range is indexes 0 to 1. The pivot is 2, which is at index 0 after the partition, so the list becomes [2, 4, 5, 6, 9]. Index 1 is bigger than 0, so the new range is index 1 only. The answer is 4.

Exercise 2: median of an odd-length list

Use quick_select to find the median of [12, 3, 9, 7, 1, 15, 4]. The median is the middle value of the sorted list.

Show the answer

Sorted, the list is 1, 3, 4, 7, 9, 12, 15. The middle item is at index 3, because the list has 7 items and indexes run from 0 to 6. The call is quick_select(a, len(a) // 2), which is quick_select(a, 3). It returns 7. In Python, len(a) / 2 would give 3.5, which is not a valid index, so use //. For an even number of items, the median is usually the average of two middle values, which needs two calls, so this lesson does not cover it.

Exercise 3: count the comparisons

Count how many comparisons quick_select makes on [1, 2, 3, 4, 5] with k equal to 0. Use the rule from the worked example: one comparison for each a[j] < pivot test.

Show the answer

The pivot is always the last item, which is the largest item of the range. Round 1 has 5 items and makes 4 comparisons. The pivot 5 lands at index 4, which is bigger than 0, so the range becomes 0 to 3. Round 2 makes 3 comparisons, then the range becomes 0 to 2. Round 3 makes 2 comparisons, and the range becomes 0 to 1. Round 4 makes 1 comparison, and the range becomes 0 to 0. The range has one item, so the loop ends. The total is 4 plus 3 plus 2 plus 1, which is 10.

Exercise 4: read the error

Run quick_select with the check on the list [5, 1, 4] and k equal to 3. Write down the error you get and say what the error tells you to change.

Show the answer

The output is a ValueError with the message "k must be from 0 to 2, got 3". The list has 3 items, so the valid indexes are 0, 1 and 2. The message says that k must be from 0 to 2, and that the value given was 3. The fix is to pass an index inside the list. If the question was about the fourth smallest value, the list has too few items to answer it.

Exercise 5: explain the random pivot

Explain in two or three plain sentences why a random pivot makes the sorted input fast. Use the numbers from the worked example.

Show the answer

With the last item as the pivot, sorted input always gives the largest item as the pivot, so each round removes only one item. A random pivot is often near the middle, so each round removes a large part of the range. In the counting program, the sorted input with the last-item pivot took 23998000 comparisons for n equal to 8000, and the random pivot took 21994 in the run shown above.

Summary

• Quickselect finds the value at a given sorted index by partitioning a range around a pivot and keeping one side.
• The index k counts from 0, and a call with an index outside the list gives a wrong answer unless the code checks it.
• The loop uses a fixed amount of memory, while a recursive version can go deep enough to hit Python's recursion limit on sorted input.
• The average cost is about a fixed multiple of n, which is linear time, and the worst cost is about n squared.
• A last-item pivot gives the worst case on sorted input. A random pivot makes that case unlikely.
• Test a selection function against sorted, with many small values and repeats.