Kth Order Statistics, in depth: quickselect with three-way partition, the expected-cost argument, heaps for streams, Floyd-Rivest, multiselect and the library calls

Define the problem precisely

Pin the definition down before writing code, because most selection bugs are definition bugs. Decide whether k is 0-indexed or 1-indexed; C++ and NumPy use 0-indexed positions, PyTorch's kthvalue uses 1-indexed k, and textbooks usually use 1-indexed. Decide what duplicates mean: with values [5, 5, 5] every rank from 1 to 3 is 5, and an algorithm that assumes distinct keys can loop or degrade. Decide what happens to NaN: it compares false with everything, so a comparison-based partition can put it anywhere and silently return a wrong answer. And decide whether a percentile is a rank (nearest-rank method) or an interpolation between two ranks; p50 of four values is either the second value or the mean of the second and third, depending on convention.

The selection family at a glance

The options trade worst-case guarantees, constants and memory. The table is the decision map this article walks through.

Quickselect with a three-way partition

Quickselect is quicksort that only recurses into one side. Choose a pivot, partition the array into values less than, equal to and greater than it, and look at the block sizes. If k falls in the less block, continue there; if in the equal block, you are done; if in the greater block, continue there with k reduced by the size of the first two blocks. The three-way partition (Dijkstra's Dutch national flag scheme) is not optional decoration: with a two-way Lomuto partition, an array of identical values makes every pivot split off a single element and the run becomes quadratic.

import random

def quickselect(a, k):
    """Return the kth smallest (0-indexed) of list a. Mutates a; copy first if needed."""
    if not 0 <= k < len(a):
        raise IndexError(k)
    lo, hi = 0, len(a) - 1
    while True:                       # iterative: no recursion-depth risk
        if lo == hi:
            return a[lo]
        pivot = a[random.randint(lo, hi)]
        lt, i, gt = lo, lo, hi        # Dutch national flag partition of a[lo..hi]
        while i <= gt:
            if a[i] < pivot:
                a[lt], a[i] = a[i], a[lt]; lt += 1; i += 1
            elif a[i] > pivot:
                a[i], a[gt] = a[gt], a[i]; gt -= 1
            else:
                i += 1
        if k < lt:
            hi = lt - 1               # answer is in the "less" block
        elif k > gt:
            lo = gt + 1               # answer is in the "greater" block
        else:
            return pivot              # k is inside the "equal" block

data = [7, 2, 5, 9, 1, 5, 8, 3, 5, 6]
assert quickselect(data[:], 5) == sorted(data)[5] == 5

After the call, the array is also partitioned around position k: everything before it is no larger and everything after it is no smaller. That side effect is often what you actually want, for example to take the k smallest items without sorting them.

Why expected work is linear

Why is the expected cost linear? Call a pivot good if it lands in the middle half of the current range, between the 25th and 75th percentile. A random pivot is good with probability one half, and a good pivot shrinks the range to at most three quarters of its size. Group the run into phases, where phase j covers ranges of size between n(3/4)^(j+1) and n(3/4)^j. The expected number of partitions in a phase is at most two (the waiting time for a coin flip), and each costs at most n(3/4)^j. Summing the geometric series gives at most 2n x 4 = 8n expected work. The bound is loose; Knuth's exact analysis gives about 2(1 + ln 2)n, roughly 3.39n, expected comparisons when k is the median, and fewer for ranks near the ends.

The worst case is still quadratic: a pivot that is always the smallest remaining element removes one value per pass. Randomization makes that astronomically unlikely for any fixed input, but if an attacker can see your random seed, or you use a fixed pivot such as the first element on already-sorted data, it happens. Introselect fixes this by counting partitions and switching to a guaranteed-linear method, such as median of medians, if progress stalls.

Streams: a heap of size k

When data arrives as a stream, or n is too large to hold, keep a max-heap of the k smallest values seen so far. Each new value is compared with the heap's top; if smaller, it replaces the top. The cost is O(n log k) time and O(k) memory, which beats quickselect whenever k is small and n cannot fit in memory. It also gives you the k smallest values, not just the kth. See Heap Operations for the underlying sift operations.

import heapq

def kth_smallest_stream(stream, k):
    """1-indexed kth smallest from an iterable, O(k) memory."""
    heap = []                          # max-heap via negation
    for x in stream:
        if len(heap) < k:
            heapq.heappush(heap, -x)
        elif x < -heap[0]:
            heapq.heapreplace(heap, -x)
    if len(heap) < k:
        raise ValueError("fewer than k items")
    return -heap[0]

Floyd-Rivest: sampling to cut comparisons

Floyd and Rivest's algorithm (1975) attacks the constant. Instead of one random pivot it draws a small random sample, selects two values from the sample that should bracket the kth element of the full array with high probability, and partitions the array into below, between and above. Usually the answer lands in the small middle band and the rest of the array is touched once. Its expected comparison count is n + min(k, n - k) + o(n), which is close to the lower bound. The cost is complexity: sample sizes and bracket offsets must be chosen carefully, and the gain matters mostly when comparisons are expensive, such as comparing long strings or calling a user-supplied comparator.

FLOYD_RIVEST_SELECT(A, k):                 # sketch, 0-indexed
    n = |A|
    if n is small: return QUICKSELECT(A, k)
    s  = sample size, about n^(2/3)
    S  = random sample of s elements of A
    g  = sqrt(s) * c                       # safety margin, c a small constant
    u  = SELECT(S, k*s/n - g)              # lower bracket
    v  = SELECT(S, k*s/n + g)              # upper bracket
    partition A into L (< u), M (u..v), R (> v)
    if |L| <= k < |L| + |M|: return SELECT(M, k - |L|)   # the common case
    else: fall back to quickselect on L or R               # rare

Several ranks at once

Asking for several ranks at once, such as p50, p90, p99 and p99.9 for a latency report, is common. Calling quickselect four times costs four linear passes. Multiselect does better: sort the requested ranks, partition once, and send each rank into the side it belongs to, recursing only where ranks remain. With r ranks this costs about O(n log r). NumPy's np.partition(a, [k1, k2, k3]) accepts a list of positions for this reason. If r approaches n, just sort.

Worked example: 10,000,000 response times, report p50, p90, p99 and p99.9. A full sort is about n log2 n, roughly 230 million comparisons. Four quickselects are roughly 4 x 3.4n, about 136 million, and many of them run on data already partly partitioned by the previous call. Multiselect does one top-level pass and then works in shrinking regions. The exact constants depend on the machine, so measure, but the ordering is reliable. For sliding windows over a stream see Moving Percentile and Median.

What libraries give you

In production, call a library. Each of these is in-place or near it, and each has a quirk worth knowing.

Two library details catch people. std::nth_element does not tell you which copy of a duplicated value you got, so attach an index if identity matters. And GPU top-k on large vocabularies, such as sampling from a language model's logits, is usually faster with sorted=False followed by a small sort of the k results.

Failure modes

• Quadratic on duplicates. Two-way partition on low-cardinality data. Fix: three-way partition.
• Quadratic on sorted input. First-element or last-element pivot. Fix: random pivot or introselect.
• Adversarial input. Predictable randomness in a public service. Fix: introselect's fallback, or median of medians.
• NaN corrupts the answer. Comparisons are false, partitions are inconsistent. Fix: filter or map NaN before selecting.
• Off-by-one rank. Mixing 0-indexed and 1-indexed APIs. Fix: name the variable for its convention and test the minimum and maximum.
• Surprise mutation. In-place selection reorders the caller's data. Fix: document it, or copy at the boundary.
• Wrong percentile convention. Dashboards disagree by one rank. Fix: choose nearest-rank or interpolation explicitly.

Trade-offs

Quickselect is the default for a single rank on in-memory data; it is simple, in place and fast in practice. The heap wins for streams and tiny k, and is the only option that needs O(k) memory. Floyd-Rivest buys fewer comparisons at the price of code you must test hard. Median of medians is slower on typical data and exists for guarantees. When the same data is queried many times for different ranks, preprocessing wins: sort once, or keep an order-statistic tree as in Kth Smallest in a BST, or, for static arrays with range queries, a merge sort tree. And if an approximate rank is enough, a quantile sketch uses far less memory than any exact method.

What to do next

1. Write down your rank convention, duplicate handling and NaN policy before choosing an algorithm.
2. Implement the iterative three-way quickselect above and test it against sorting on random, sorted, constant and two-value arrays.
3. Add the heap version for streams and check it uses O(k) memory.
4. Replace hand-written selection in production code with nth_element, np.partition or torch.topk, and note their indexing.
5. Batch percentile reports into one multiselect call.
6. If inputs can be adversarial, confirm your library falls back to a worst-case linear method; read Median of Medians and Quicksort for the partition details.
7. Benchmark on your own data sizes before choosing Floyd-Rivest or a sketch.