Searching and sorting / Sorting / 044 · Quickselect (kth element)

Quickselect (kth element)

Partition around a pivot; recurse only into the side holding k.

The invariant

After partition, the pivot is in its final sorted position.

01The intuition

Start with what you would do by hand

The list 3, 2, 1, 5, 6, 4. What is the 2nd largest number? Sorting gives 1, 2, 3, 4, 5, 6, and the 2nd largest is 5. Sorting is easy but it does more than asked: it puts all six numbers in order when we only care about one position.

Sorting n numbers costs about n log2 n comparisons, around 20 million for a million numbers. Can we find one position with about n comparisons?

Here is a cheaper by-hand move. Pick any number, say 4, and split the others into "smaller than 4" and "at least 4": {3, 2, 1} and {5, 6}. Now 4 is in its final sorted position, index 3, because exactly three numbers are smaller. The 2nd largest is at sorted index 4, which is to the right of 4. So we can throw away the left group entirely and continue with {5, 6}. We never sorted the left group.

Build the structure, one piece at a time

The structure is the array itself, rearranged in place, plus a range [lo, hi] that shrinks.

Piece 1: target index. The k-th largest of n numbers sits at sorted index n − k. Here k = 2 and n = 6, so the target index is 4.

Piece 2: partition. Choose a pivot (the last element, 4). Keep a "store" index starting at lo. Walk the range; every time you see a value smaller than the pivot, swap it into the store position and advance the store. Finally swap the pivot into the store position. Everything left of the store is smaller than the pivot; the pivot is at its final sorted position.

StepArrayStore
start, pivot 4 at the end3, 2, 1, 5, 6, 40
3 < 4: swap into slot 03, 2, 1, 5, 6, 41
2 < 4: swap into slot 13, 2, 1, 5, 6, 42
1 < 4: swap into slot 23, 2, 1, 5, 6, 43
5, 6 not smaller: leave3, 2, 1, 5, 6, 43
pivot to slot 33, 2, 1, 4, 6, 5pivot at 3

Piece 3: recurse into one side. The pivot landed at 3, the target is 4, so look only in [4, 5]. Pivot 5: 6 is not smaller, so the pivot lands at index 4. That is the target: the 2nd largest is 5.

Why it is correct

The invariant: the target index lies inside [lo, hi], and every element left of lo is smaller than every element inside the range, and every element right of hi is larger or equal.

After a partition, the pivot sits at index p with everything smaller to its left and everything else to its right, so the pivot is exactly the element that belongs at sorted index p. If p equals the target, the pivot is the answer. If the target is less than p, the answer is among the smaller elements, which are all in [lo, p − 1]; if greater, in [p + 1, hi]. Discarding the other side never discards the answer.

after partition: array[i] < pivot for i < p, array[i] ≥ pivot for i > p, so pivot = sorted[p]

Why it is fast

Each partition walks the current range once. If the pivot splits the range roughly in half each time, the ranges are n, n/2, n/4, …

n + n/2 + n/4 + … < 2n, so O(n) on average

Compare with quicksort, which must recurse into both sides and so costs n log n. A bad pivot (always the smallest) gives ranges n, n − 1, … and O(n2); choosing the pivot at random makes that vanishingly unlikely. For a million numbers, about two million comparisons instead of twenty million.

How to recognise it in an interview

The cues are "k-th largest", "k-th smallest", "median", "top k" when the order among the top k does not matter, and especially "can you do better than sorting". Say: "I partition around a pivot and recurse only into the side that contains the k-th index, which is O(n) average."

Kth Largest Element in an Array is the example above. Top K Frequent Elements runs quickselect over the (value, count) pairs comparing counts. K Closest Points to Origin partitions points by squared distance. When a stream is involved, as in Kth Largest Element in a Stream, a heap is the right tool instead, because quickselect needs the whole array.

02The idea in one line

Partition around a pivot; recurse only into the side holding k.

Reach for it when: Kth largest, top-k without a heap, median.

Cost: O(n) average time.

03The walk-through, step by step

The video works one example: Kth largest (k=2) in [3,2,1,5,6,4]. These are its captions in order, so you can follow without the video or check a step you missed.

  1. Partition around a pivot; recurse only into the side that holds index 4
  2. Range [0, 5], pivot = last = 4
  3. 3 < 4: swap into slot 0, store = 1
  4. 2 < 4: swap into slot 1, store = 2
  5. 1 < 4: swap into slot 2, store = 3
  6. 5 ≥ 4: leave it on the right
  7. 6 ≥ 4: leave it on the right
  8. Place the pivot at index 3: everything left is smaller
  9. 3 < 4: the answer is on the right
  10. Range [4, 5], pivot = last = 5
  11. 6 ≥ 5: leave it on the right
  12. Place the pivot at index 4: everything left is smaller
  13. Index 4 is the target: the 2nd largest is 5
  14. One side per round, halving on average: O(n) expected

04Asked at Google

5 Google-tagged LeetCode problems use this pattern. Frequency is LeetCode's Google-tag score (higher means reported more often); "last 6 months" marks problems reported in the most recent window of the export.

ProblemDifficultyFrequencyReported
Top K Frequent ElementsMedium59last 6 months
Kth Largest Element in an ArrayMedium55last 6 months
Find Median from Data StreamHard50last 6 months
Kth Largest Element in a StreamEasy33last 6 months
K Closest Points to OriginMedium27

Other patterns in Sorting: Custom comparator / key sort, Cyclic sort (1..n placement), Merge sort with counting, Bucket / counting sort.