Median of a large unsorted array in linear time

The median is a selection problem, not a sorting problem. For odd n we need the element of rank n // 2. For even n we need the elements of ranks n // 2 - 1 and n // 2. We do not need the rest of the array in order.

Brute force. Sort a copy and read the middle one or two values. That is O(n log n) time and O(n) space. Say it first: it is the correct baseline and the bar the optimal answer must beat.

Heaps. A max-heap of size n // 2 + 1 (or the two-heap method from the streaming version) gives O(n log n) too. It is the right tool when values arrive one at a time, but here the whole array is already in hand.

Optimal: quickselect (expected O(n)).

1. Pick a random pivot in the active range [lo, hi].
2. Partition the range three ways: < pivot, == pivot, > pivot. Call the boundaries of the equal block lt and gt.
3. If the target rank k is inside [lt, gt], the answer is the pivot. If k < lt, continue only on the left part. If k > gt, continue only on the right part.

Each round throws away one side. With a random pivot the active range shrinks by a constant factor on average, so the total work is n + 3n/4 + (3/4)^2 n + ... = O(n) expected.

Why three-way partitioning. With a two-way partition, an array full of duplicates (for example all zeros) makes every pivot split badly and the run becomes O(n^2). The equal block absorbs all copies of the pivot in one pass, so duplicates make the search end sooner, not later.

Even length without a second select. After quickselect(a, mid) returns, the array satisfies a[:mid] <= a[mid] <= a[mid+1:]. So the lower middle value is just max(a[:mid]), one linear scan. A second quickselect would also work, but the scan is simpler and cheaper.

Invariants.

• At the start of each round, the target rank k is inside [lo, hi], every value left of lo is <= every value in the range, and every value right of hi is >= it.
• After _partition3, a[lo..lt-1] < pivot, a[lt..gt] == pivot, a[gt+1..hi] > pivot.

Details to state up front. Return a float, because the even case is an average. Raise on an empty array. Work on a copy unless the caller says the input may be reordered (for a truly huge array, ask: mutating in place saves O(n) memory). Python integers do not overflow, but in Java or C++ compute the average as lo + (hi - lo) / 2.0 or in a wider type.

import random
from typing import Callable, Iterable


def _partition3(a: list[int], lo: int, hi: int, pivot: int) -> tuple[int, int]:
    """Dutch-flag partition of a[lo..hi] around pivot.

    Returns (lt, gt): a[lo..lt-1] < pivot, a[lt..gt] == pivot, a[gt+1..hi] > pivot.
    """
    lt, i, gt = lo, 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
    return lt, gt


def quickselect(a: list[int], k: int) -> int:
    """Return the k-th smallest (0-indexed) value. Reorders a in place so that
    a[:k] <= a[k] <= a[k+1:]."""
    lo, hi = 0, len(a) - 1
    while True:
        pivot = a[random.randint(lo, hi)]
        lt, gt = _partition3(a, lo, hi, pivot)
        if k < lt:
            hi = lt - 1
        elif k > gt:
            lo = gt + 1
        else:
            return pivot  # a[lt..gt] all equal pivot, and k is inside


def find_median(nums: list[int]) -> float:
    """Median of an unsorted list in expected O(n) time."""
    n = len(nums)
    if n == 0:
        raise ValueError("median of empty array")
    a = list(nums)  # drop this copy if the caller allows mutation
    mid = n // 2
    upper = quickselect(a, mid)
    if n % 2 == 1:
        return float(upper)
    lower = max(a[:mid])  # after selection, everything left of mid is <= upper
    return (lower + upper) / 2


# Follow-up: the data does not fit in memory. We can only re-scan it.
def kth_by_value_search(scan: Callable[[], Iterable[int]], k: int) -> int:
    """k-th smallest (0-indexed) using O(1) memory and O(log(range)) passes."""
    lo = hi = None
    count = 0
    for x in scan():
        lo = x if lo is None else min(lo, x)
        hi = x if hi is None else max(hi, x)
        count += 1
    if count == 0 or not 0 <= k < count:
        raise ValueError("k out of range")
    # Smallest v with #{x <= v} > k is the k-th smallest.
    while lo < hi:
        mid = (lo + hi) // 2
        if sum(1 for x in scan() if x <= mid) > k:
            hi = mid
        else:
            lo = mid + 1
    return lo


def median_by_value_search(scan: Callable[[], Iterable[int]], n: int) -> float:
    upper = kth_by_value_search(scan, n // 2)
    if n % 2 == 1:
        return float(upper)
    return (kth_by_value_search(scan, n // 2 - 1) + upper) / 2

• find_median: expected O(n) time. The worst case is O(n^2) if every random pivot is extreme, which has vanishing probability. Space is O(n) for the defensive copy and O(1) extra beyond it (the loop is iterative, so no recursion stack). Without the copy it is O(1) extra.
• The even case adds one O(n) scan for max(a[:mid]); the slice itself is O(n) memory, which a plain loop over indices avoids if memory is tight.
• kth_by_value_search: O(n log R) time, where R is max - min, because each pass counts over the whole input and the value range halves each pass. That is about 32 passes for 32-bit integers. Extra memory is O(1).
• median_by_value_search: at most two such searches.

• [3, 1, 2] -> 2.0 (odd length).
• [7, 1, 4, 10] -> 5.5 (even length, average of 4 and 7).
• [5] -> 5.0 (single element).
• [2, 2, 2, 2] -> 2.0 (all equal; three-way partition finishes in one round).
• [-5, -1, -3, -2] -> -2.5 (negatives, non-integer result).
• [1, 1, 1, 9, 9, 9] -> 5.0 (the two middle values sit in different duplicate blocks).
• [] -> raises ValueError.
• [0] * 100_000 -> 0.0 quickly (the duplicate case that breaks a two-way partition).
• 200,001 random values in [-10^9, 10^9] -> same as statistics.median.
• The input list is unchanged after the call.

1. Guaranteed linear time. Replace the random pivot with median of medians: split into groups of 5, take each group's median, recursively select the median of those as the pivot. That pivot is guaranteed to discard at least about 30% of the range each round, so the worst case is O(n). Mention that the constant factor is large and that production code (for example introselect in C++ nth_element) starts with random or median-of-3 pivots and falls back to median of medians only when progress is poor.
2. Numbers arrive as a stream (LeetCode 295). Keep a max-heap for the lower half and a min-heap for the upper half, with sizes equal or the lower half one larger. add is O(log n), median is O(1). If the stream also expires old values (sliding window), use lazy deletion or two sorted containers.
3. Too large for memory, but the data can be re-read. Binary search on the value, not the index: for a candidate v, one pass counts how many values are <= v. That is kth_by_value_search above, O(log R) passes with O(1) memory. A faster variant reads once into a histogram of the top 16 bits, finds the bucket that holds rank k, then makes a second pass on that bucket only.
4. Data sharded across many machines. Use the same value-space binary search: the coordinator broadcasts a candidate, each shard returns its local count <= v, and the coordinator sums them. That costs O(log R) rounds and tiny messages. An alternative is distributed quickselect: shards report counts below and above a shared pivot, and each round discards a side everywhere. If an approximate answer is fine, merge per-shard sketches (t-digest, KLL) in one round.
5. Small value range, or the k-th smallest in general. If values are bounded (ages, ratings), counting sort gives O(n + R) exactly. The same quickselect returns any rank k, so the 90th percentile or the top-k elements (everything right of the selected index) come with no extra code.

Mid-level (L4)

• States the sort-based O(n log n) answer, then reaches quickselect with a little guidance.
• Writes a working partition with correct boundaries and handles odd and even lengths.
• Gives the expected O(n) bound and knows the worst case is quadratic.

Senior (L5)

• Goes to quickselect directly with a random pivot and explains the geometric-series argument for expected O(n).
• Uses a three-way partition, or at least spots the all-duplicates trap without prompting.
• Gets the even case with one select plus a scan, raises the copy-versus-mutate trade-off, and tests edge cases (empty, one element, duplicates, negatives).
• Answers the streaming follow-up with two heaps.

Staff (L6+)

• Explains median of medians and why production code uses introselect instead.
• Handles out-of-memory and distributed data with value-space binary search or histogram passes, and reasons about passes, rounds and network cost rather than only Big-O.
• Knows when an approximate median from a mergeable sketch is the right product answer, and asks what accuracy the caller needs.