Selection Algorithms, in depth: comparison lower bounds, sorted inputs, weighted medians, radix select and distributed selection

The selection family

Formally, given n items with a total order and an integer k between 1 and n, return the item that would sit at position k if the items were sorted. Ties matter: with duplicates there may be many positions holding the same value, and a correct implementation returns the value, not a particular index, unless it promises stability. Variants change the input, the output or the cost model:

• Single rank in an unsorted array: quickselect and its relatives, covered in Kth order statistics and, for the worst-case guarantee, median of medians.
• Top k, sorted: select the k-th, then sort the prefix, or keep a bounded heap; see partial sort.
• Sorted or semi-sorted inputs: two sorted arrays, a row-and-column sorted matrix, a merge of many sorted runs. Here selection drops below linear time.
• Weighted ranks: each item carries a weight and you want the item where the cumulative weight crosses half (or any fraction) of the total.
• Fixed-width keys on parallel hardware: radix select, where you never compare two keys at all.
• Approximate ranks over streams: sketches such as t-digest trade exactness for one pass and bounded memory.

What comparisons can and cannot do

Before optimising anything it helps to know what is possible. In the comparison model, where the only way to learn about keys is to compare two of them, several selection problems have exact answers.

Two practical lessons fall out of this table. First, finding the minimum and maximum together in pairs saves a quarter of the comparisons, which is worth doing in hot loops that compute ranges for histograms or normalisation. Second, the gap between the median bound and what practical deterministic code does is large: the textbook median-of-medians runs well above 5n comparisons in the worst case, which is why libraries use randomised or introspective pivots and keep the deterministic algorithm only as a fallback.

The pairing version is a ten-line loop: order each pair with one comparison, then compare only the smaller with the running minimum and the larger with the maximum.

Selecting from sorted inputs

When the input is already sorted in pieces, you can select without touching most elements. The classic case is the k-th smallest of two sorted arrays A and B. Take i elements from A and k - i from B; the split is right exactly when the last taken element of each array is no larger than the first untaken element of the other. That condition is monotone in i, so binary search finds it in O(log min(m, n)) steps:

def kth_of_two(A, B, k):
    """k is 1-based. A and B are sorted ascending."""
    if len(A) > len(B):
        A, B = B, A
    lo, hi = max(0, k - len(B)), min(k, len(A))
    NEG, POS = float("-inf"), float("inf")
    while lo <= hi:
        i = (lo + hi) // 2            # take i from A
        j = k - i                     # and j from B
        a_left  = A[i - 1] if i > 0 else NEG
        a_right = A[i]     if i < len(A) else POS
        b_left  = B[j - 1] if j > 0 else NEG
        b_right = B[j]     if j < len(B) else POS
        if a_left <= b_right and b_left <= a_right:
            return max(a_left, b_left)
        if a_left > b_right:
            hi = i - 1                # took too many from A
        else:
            lo = i + 1                # took too few from A
    raise ValueError("k out of range")

For many sorted runs, the same idea generalises to a search on value instead of index. Guess a value v, count how many elements are at most v in each run with a binary search, and adjust v. For an n by n matrix whose rows and columns are both sorted, the count costs O(n) with a staircase walk from the bottom-left corner, so the whole search costs O(n log V) where V is the value range. The lesson is general: whenever you can count elements below a threshold faster than you can enumerate them, binary search on the answer turns counting into selection. Range structures such as the merge sort tree are built on exactly that move.

The weighted median

In a weighted median each item carries a non-negative weight, and the answer is the item where the running weight first reaches half the total. It minimises the sum of weighted absolute distances, which is why it appears in facility location, robust regression and in the pivot step of some parallel selection algorithms. Sorting gives O(n log n). A selection-style recursion gives expected O(n): pick a pivot, sum the weight on each side, and recurse only into the side that holds the crossing point, carrying the weight you have already accounted for.

import random

def weighted_median(items, target_frac=0.5):
    """items: list of (value, weight) with weight >= 0. Returns a value."""
    total = sum(w for _, w in items)
    target = target_frac * total
    below = 0.0                                   # weight already known to be smaller
    while True:
        pv = random.choice(items)[0]
        lo = [(v, w) for v, w in items if v < pv]
        eq = sum(w for v, w in items if v == pv)
        w_lo = sum(w for _, w in lo)
        if below + w_lo >= target:                # crossing point is left of pivot
            items = lo
        elif below + w_lo + eq >= target:         # pivot itself crosses the target
            return pv
        else:
            below += w_lo + eq
            items = [(v, w) for v, w in items if v > pv]

The three-way split around the pivot matters: with heavy duplicates, a two-way split can recurse forever on an array of equal values. Decide whether you want the lower or upper weighted median and test both.

Radix select: no comparisons at all

Comparison selection is inherently sequential within a partition pass and branchy, which suits CPUs and wastes GPUs. Radix select avoids comparisons entirely. Treat each key as a fixed-width unsigned integer, look at its most significant b bits, and build a histogram of 2^b buckets. A prefix sum over the histogram tells you which bucket contains rank k and how many elements sit in lower buckets. Keep only that bucket, subtract the lower count from k, and repeat on the next b bits. After w / b passes over a shrinking candidate set you have the exact key.

Floating-point keys need one trick: flip the bits so unsigned order matches float order. For a non-negative float set the sign bit; for a negative float invert every bit. NaNs must be filtered or given an explicit rank first.

import struct

def float_key(x):
    u = struct.unpack("<I", struct.pack("<f", x))[0]
    return u ^ 0xFFFFFFFF if u & 0x80000000 else u | 0x80000000

def radix_select(values, k, bits=8):
    """k is 1-based rank of the smallest. 32-bit float keys, 4 passes of 8 bits."""
    cand = [(float_key(v), v) for v in values]
    for shift in range(32 - bits, -1, -bits):
        hist = [0] * (1 << bits)
        for key, _ in cand:
            hist[(key >> shift) & ((1 << bits) - 1)] += 1
        seen = 0
        for digit, count in enumerate(hist):        # prefix sum finds the bucket
            if seen + count >= k:
                break
            seen += count
        k -= seen
        cand = [(key, v) for key, v in cand if (key >> shift) & ((1 << bits) - 1) == digit]
    return cand[0][1]

On a GPU, each pass is a parallel histogram in shared memory followed by a tiny prefix sum, and the work is perfectly regular: no data-dependent branches, no pivot luck. PyTorch's CUDA top-k is built on a radix-select kernel for this reason, and it is why selecting the top 50 logits from a 150,000-entry vocabulary costs almost nothing compared with sorting it.

Worked example: radix select by hand

Take eight 8-bit keys and find the 3rd smallest using 4-bit digits: 0x5A, 0x13, 0x57, 0xC2, 0x1F, 0x50, 0x9E, 0x11. In decimal those are 90, 19, 87, 194, 31, 80, 158 and 17, so the answer should be 31.

1. Pass 1, high nibble. Histogram: digit 1 holds 3 keys (0x13, 0x1F, 0x11), digit 5 holds 3, digit 9 holds 1, digit C holds 1. Running count reaches 3 at digit 1, and 3 is at least k = 3, so keep digit 1. Nothing sat in lower buckets, so k stays 3.
2. Pass 2, low nibble of the three survivors: 0x13 has 3, 0x1F has F, 0x11 has 1. Digits in order are 1, 3, F with one key each. The running count reaches 3 at F.
3. The survivor is 0x1F, which is 31. Two passes, eight plus three keys examined, and not a single key-to-key comparison.

External and distributed selection

When data lives on many machines or on disk, the cost that matters is passes over the data and rounds of communication, not comparisons. Distributed selection runs pivot rounds: a coordinator proposes a pivot, every shard counts how many of its local candidates are below, equal to and above it, the coordinator sums the counts, decides which side holds rank k, and every shard discards the other side locally. Only three integers per shard cross the network per round.

def distributed_select(shards, k):
    """shards: objects exposing sample(), count(p) -> (lt, eq, gt), keep(side, p)."""
    while True:
        p = weighted_median_of(s.sample() for s in shards)    # good pivot, cheap to send
        lt = eq = 0
        counts = [s.count(p) for s in shards]                 # one RPC per shard
        for a, b, _ in counts:
            lt += a; eq += b
        if k <= lt:
            for s in shards: s.keep("lt", p)
        elif k <= lt + eq:
            return p
        else:
            k -= lt + eq
            for s in shards: s.keep("gt", p)

Choosing the pivot as the weighted median of shard medians, weighted by shard size, keeps every round discarding a constant fraction, so you need O(log n) rounds. A sampling variant does better: draw a few thousand random elements, pick two pivots that bracket rank k with high probability (the Floyd-Rivest idea applied across machines), count once, and ship only the elements between the pivots to one node. That is usually two rounds in total. The same two-pass plan works for a file larger than memory: one streaming pass with a reservoir sample to choose bracketing pivots, a second pass that counts below the lower pivot and collects the band between them, then an in-memory select on the band.

Failure modes

• Adversarial or sorted input with naive pivots. First-element pivots turn quickselect quadratic on sorted data. Use random or median-of-three pivots with an introselect depth limit.
• Duplicates. Two-way partitions on many equal keys fail to shrink. Always partition three ways or count equals explicitly, as the code above does.
• NaN and mixed types. NaN breaks the total order and can make a selection loop spin or return garbage. Filter or rank NaNs before selecting.
• Off-by-one ranks. Mixing 0-based and 1-based k is the most common bug in selection code. Property-test against sorted(a)[k - 1] on random small inputs.
• Radix skew and wide keys. 64-bit keys with 8-bit digits need eight passes; if candidates do not shrink because keys share prefixes, memory traffic stays high.

Trade-offs

Comparison methods are general (any ordered type, any comparator) and fast on CPUs, but their running time depends on pivot luck unless you pay for a deterministic fallback. Radix select is predictable and parallel, but only for keys that map to fixed-width integers, and its memory traffic is several passes. Binary search on the answer is the cheapest when the input is already organised, and useless when it is not. Distributed exact selection costs rounds and stragglers; sketches cost accuracy. Choose by input shape first, then by guarantee, then by constant factors.

What to do next

1. Find every place your code sorts only to read one rank or the first k items, and replace it with a selection call such as numpy.partition or std::nth_element.
2. Write a property test that compares your selector against sorted(a)[k - 1] on random inputs with duplicates, NaN handling and k at both ends.
3. If you compute a min and max in the same loop, switch to the pairing version.
4. If inputs arrive as sorted runs, try binary search on the value with a counting function before reaching for a merge.
5. For GPU top-k, check that your framework's top-k is used instead of a full sort.
6. For sharded data, decide whether exact rank is a hard requirement; if not, deploy a mergeable quantile sketch per shard.