39. Kth Largest Eval Score via Quickselect

Problem

A batch eval job scores a large pool of model checkpoints on a held-out set. To pick a cutoff for the next promotion round, an engineer wants the kth-largest score out of the batch, and this runs on every batch, so a full O(n log n) sort just to read off one rank is wasteful.

Given the unsorted scores, return the kth largest without fully sorting. k is 1-indexed (k=1 means the single largest score). This must run in average O(n) time using quickselect (partition-based selection), not a full sort.

Source: src/39_kth_largest_score_quickselect.py

def kth_largest_score(scores: list[float], k: int) -> float:
    ...

Examples:
>>> kth_largest_score([3.1, 1.5, 9.2, 4.4, 7.0], 2)
7.0

>>> kth_largest_score([5.0], 1)
5.0

Step-by-Step Approach

1. Translate "kth largest" into an ascending-sorted-index target: the kth largest is the element that would sit at index n - k if the array were fully sorted ascending.
2. Work on a copy of the input array, since partitioning rearranges elements in place and the caller's list should not be mutated as a side effect.
3. Pick a pivot (a random index, to avoid the O(n²) worst case on adversarial or already-sorted input) and partition the array around it, so everything smaller ends up to its left and everything larger to its right.
4. After partitioning, the pivot sits at its final sorted position, pivot_idx. Compare it to the target index.
5. If pivot_idx == target, you're done — return that value. If pivot_idx < target, the answer lies to the right, so continue only on the right sub-range. Otherwise continue only on the left sub-range.
6. Because each step only recurses into one side (never both, unlike quicksort), the expected work halves geometrically, giving average O(n) instead of O(n log n).
7. Handle the base case where the search range collapses to a single index (lo == hi) — that element is the answer regardless of further partitioning.

The key insight is that quickselect throws away the partition it doesn't need instead of recursing into both, which is exactly what turns an O(n log n) sort into an average O(n) selection.

Reference solution

import random


def _partition(arr, lo, hi, pivot_idx):
    pivot_val = arr[pivot_idx]
    # move pivot to the end out of the way
    arr[pivot_idx], arr[hi] = arr[hi], arr[pivot_idx]
    store = lo
    for i in range(lo, hi):
        if arr[i] < pivot_val:
            arr[store], arr[i] = arr[i], arr[store]
            store += 1
    # move pivot into its final sorted position
    arr[store], arr[hi] = arr[hi], arr[store]
    return store


def kth_largest_score(scores: list[float], k: int) -> float:
    arr = list(scores)  # work on a copy, never mutate the caller's list
    n = len(arr)
    # kth largest == the element at ascending-sorted index (n - k)
    target = n - k
    lo, hi = 0, n - 1
    while True:
        if lo == hi:
            return arr[lo]
        # random pivot avoids the O(n^2) worst case on adversarial/sorted input
        pivot_idx = random.randint(lo, hi)
        pivot_idx = _partition(arr, lo, hi, pivot_idx)
        if pivot_idx == target:
            return arr[pivot_idx]
        elif pivot_idx < target:
            lo = pivot_idx + 1
        else:
            hi = pivot_idx - 1

Key Functions & Tricks

• random.randint(lo, hi) — randomized pivot choice, defeats adversarial inputs that would make a fixed pivot choice O(n²)
• Lomuto partition — single left-to-right scan that places the pivot at its final sorted index in O(hi - lo) time
• target = n - k — converts "kth largest" into the equivalent ascending-order index to search for
• One-sided recursion — unlike quicksort, only the sub-range containing the target is ever explored, which is what gives average O(n) instead of O(n log n)
• arr = list(scores) — defensive copy so the in-place partitioning doesn't mutate the caller's input list
• lo == hi base case — a single-element range is trivially its own answer

How to Recognize This Pattern