Quickselect

Problem

Given an unsorted array, find the element that would sit at position k if the array were sorted: the k-th smallest, the k-th largest, the median, or the k nearest items to a point. You are allowed to rearrange the array in place, you need only that single position (or the group of items before it), and the input is fixed rather than a stream. Sorting does too much, and a heap of size k costs O(n log k).

Intuition

Suppose you want the 6th shortest of eight people. Pick anyone as the pivot, say Sam, and ask everybody to stand on Sam's left if shorter, and on Sam's right otherwise. Sam is now in the exact place he would take in a fully sorted line, and you can count his position. If it is the 6th place, you are done. If Sam is 3rd, the 6th shortest must be among the people on the right, so everyone on the left can go home. Repeat inside the smaller crowd, asking for the 3rd shortest of those who remain.

A fair pivot cuts the crowd roughly in half, so the crowds are n, then about n/2, then about n/4, and the total work is about 2n. A terrible pivot, such as always the tallest person, removes only one person per round and the work becomes n + (n - 1) + (n - 2) + ... = O(n²). That is why a random pivot matters.

A partition does not sort either side. It only guarantees that the pivot is in its final place with the smaller values before it. The diagram scrolls sideways on a narrow screen.

Mechanics

The invariant

The live range [lo, hi] always contains sorted position target, and every value left of lo is no larger than every value in the range while every value right of hi is no smaller. Partitioning the range around a pivot puts the pivot at an index p in its final place within the range. If p equals the target, the pivot is the answer. If the target is less than p, the answer is in [lo, p - 1]; otherwise it is in [p + 1, hi]. Either way the invariant holds on the smaller range, and since the range strictly shrinks, the loop ends.

The Lomuto partition

Take the last item as the pivot. Keep a boundary just after the region of values known to be smaller than the pivot. Scan the range: whenever you meet a value smaller than the pivot, swap it into the boundary slot and move the boundary right. At the end swap the pivot into the boundary slot. All values left of it are smaller; all values to its right are at least as large. One pass, O(range) comparisons, O(1) extra space.

Expected and worst cost

With a pivot that lands anywhere in the middle half, the range keeps at most three quarters of its items, so the total comparisons are at most n(1 + 3/4 + (3/4)² + ...) = 4n, which is O(n). Choosing the pivot at random makes a bad split (outside the middle half) happen with probability one half in each round, and the expected total stays O(n) for every input. Without randomization, an adversarial or already-sorted input with the last item as pivot removes one item per round and costs n(n - 1)/2 comparisons.

Green cells are the live range going into the round, amber is where that round's pivot landed. The range shrinks until the pivot lands on the target index.

Worked example

Find the 3rd largest of 9, 4, 7, 1, 8, 3, 6, 2, which is the value at ascending index 5 (8 - 3). Each row is one partition of the live range.

The answer is 7, found after 5 partitions and 18 comparisons, without ever sorting the whole array. The search can stop with parts of the array unsorted; here the last partitions happened to leave it in order. Round 5 placed the pivot at index 6, which is past the target 5, so the live range shrank to the single index 5, whose value was the answer.

The bad case. On the already sorted array 1 to 8, looking for the smallest, the last-item pivot is the largest each time. That costs 7 rounds and 28 comparisons for 8 items, which is 8 × 7 / 2: quadratic. A random pivot avoids it.

Implementation

import random


def find_kth_largest(nums, k):
    nums = list(nums)                       # keep the caller's array intact
    target = len(nums) - k
    lo, hi = 0, len(nums) - 1
    while lo < hi:
        pick = random.randint(lo, hi)
        nums[pick], nums[hi] = nums[hi], nums[pick]
        pivot = nums[hi]
        boundary = lo
        for j in range(lo, hi):
            if nums[j] < pivot:
                nums[boundary], nums[j] = nums[j], nums[boundary]
                boundary += 1
        nums[boundary], nums[hi] = nums[hi], nums[boundary]
        if boundary == target:
            return nums[boundary]
        if target < boundary:
            hi = boundary - 1
        else:
            lo = boundary + 1
    return nums[target]


def k_closest(points, k):
    def dist(p):
        return p[0] * p[0] + p[1] * p[1]

    lo, hi = 0, len(points) - 1
    while lo < hi:
        pick = random.randint(lo, hi)
        points[pick], points[hi] = points[hi], points[pick]
        pivot = dist(points[hi])
        boundary = lo
        for j in range(lo, hi):
            if dist(points[j]) < pivot:
                points[boundary], points[j] = points[j], points[boundary]
                boundary += 1
        points[boundary], points[hi] = points[hi], points[boundary]
        if boundary == k:
            break
        if k < boundary:
            hi = boundary - 1
        else:
            lo = boundary + 1
    return points[:k]

Swapping a random item to the end first is what randomizes the pivot while keeping the same partition code. The loop is iterative on purpose: only one side is ever followed, so no recursion is needed.

Variations and edge cases

• Many duplicates. The partition above puts equal values on the right, so an array of equal items degrades to quadratic time. Use a three-way partition (smaller, equal, larger) and stop as soon as the target falls in the equal block.
• Median-of-medians. Choosing the pivot as the median of group medians guarantees a good split and worst-case O(n), but with a large constant, so it is rarely worth it in practice.
• The k best as a set. After placing index k, the first k entries are the k smallest in no particular order. This is the quickselect answer for "k closest" problems.
• Index mix-ups. The k-th largest sits at ascending index n - k, and the k-th smallest at k - 1 when counting from one. Write the target on one line and use it everywhere.
• Edge cases: a single element, k equal to 1 or n, k out of range (validate first), all values equal, and callers who do not expect their array to be reordered (copy it, or say so).

Complexity and tradeoffs

• Time: expected O(n) with a random pivot, O(n²) worst case, with the worst case made very unlikely by randomization. Median-of-medians makes it O(n) always.
• Space: O(1) extra when you work in place on the array, O(n) if you copy it first.
• Baseline. A heap of size k costs O(n log k) with no mutation and handles streams. The sort is O(n log n) and the shortest to write. Quickselect wins when n is large and you may rearrange the array.
• When it fails. Streams, read-only or shared inputs, hard latency limits where an unlucky pivot would matter, and repeated queries for different k: sort once instead, then every query is O(1).

Problems

The home problem is the classic k-th largest of an array, which can be solved by this page's partition or by the heap from the previous pages. The sort and closest-points problems listed under "also uses" show quicksort-style partitioning and the k-nearest variant.

Medium

• 215. Kth Largest Element in an Array Grind 169 NeetCode 150 NeetCode 250
  Either keep a min-heap of the k largest values seen, or partition around a pivot until the pivot lands at index n-k.

Also uses this pattern

These problems have a different home page, but this idea is part of their solution.

• 912. Sort an Array Medium
• 973. K Closest Points to Origin Medium

Connections

This page reuses the swap-based scanning of the two-pointers page, and is the in-place alternative to the heap of the top-k page. Comparing the two is a favorite interview follow-up: the heap is streaming-friendly and has a guaranteed O(n log k), quickselect has a better average and no extra memory. Discarding part of the search space each step is the same move as binary search, with a partition standing in for the midpoint.

Check yourself

References

• Hoare, C. A. R., "Algorithm 65: Find", Communications of the ACM, 1961, the original selection algorithm.
• Cormen, Leiserson, Rivest and Stein, Introduction to Algorithms, chapter 7 (quicksort) and chapter 9 (randomized selection and the worst-case linear-time algorithm).
• Blum, Floyd, Pratt, Rivest and Tarjan, "Time bounds for selection", 1973, the median-of-medians result.
• Sedgewick and Wayne, Algorithms, section 2.3 (quicksort and selection, three-way partitioning).
• Skiena, The Algorithm Design Manual: the quicksort section of the sorting chapter, and the catalog entry on median and selection.
• The NeetCode roadmap and the Blind 75, Grind 169, NeetCode 150 and NeetCode 250 lists, used as the index of topics and problems.