28. Kth Largest Rerank Score

Problem

After Fin's reranker scores a batch of candidate passages, a monitoring job occasionally needs to pull the kth-best score from the batch, e.g. to check where the score cutoff for the top-3 shown to the model would fall. This runs on every single query, so fully sorting the batch just to read off one order statistic is wasteful.

Given the unsorted scores and an integer k, return the kth largest score in the batch WITHOUT sorting the whole batch.

Source: src/28_kth_largest_rerank_score.py

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

kth_largest([3.0, 2.0, 1.0, 5.0, 6.0, 4.0], 2)
# -> 5.0   (sorted desc: 6, 5, 4, 3, 2, 1 -- 2nd largest is 5)

kth_largest([3.0, 2.0, 3.0, 1.0, 2.0, 4.0, 5.0, 5.0, 6.0], 4)
# -> 4.0

Step-by-Step Approach

1. Copy scores into a working array arr so the input isn't mutated, and compute target_index = len(arr) - k — the index the kth largest value would land at if arr were sorted ascending.
2. Write a partition(left, right, pivot_index) helper: move the pivot value to the end, then walk left to right swapping any element smaller than the pivot into a growing "less than pivot" region, and finally swap the pivot into its correct sorted position.
3. Set left = 0, right = len(arr) - 1 and loop while left < right.
4. Each iteration, pick a random index in [left, right] as the pivot (randomizing avoids worst-case behavior on adversarial or already-sorted input) and partition around it.
5. If the pivot lands exactly at target_index, that value is the answer — stop. If the pivot's final index is less than target_index, the answer is to its right, so search [pivot_index + 1, right]. Otherwise search [left, pivot_index - 1].
6. Once the loop ends (or breaks early), return arr[target_index].

The key insight is that finding the kth order statistic doesn't require a total order over the whole array — each partition step only needs to narrow down which side the target index falls on, discarding the other side entirely. This gives average O(n) time versus O(n log n) for a full sort.

Reference solution

import random


def kth_largest(scores: list[float], k: int) -> float:
    # quickselect with a random pivot: average O(n) time, O(n) space for the working copy
    # shallow copy, caller's list untouched
    arr = list(scores)
    # index of the kth largest once arr is sorted ascending
    target_index = len(arr) - k

    def partition(left: int, right: int, pivot_index: int) -> int:
        pivot_value = arr[pivot_index]
        # swap pivot to end (no temp var)
        arr[pivot_index], arr[right] = arr[right], arr[pivot_index]
        store_index = left
        for i in range(left, right):
            if arr[i] < pivot_value:
                # move smaller element into place
                arr[store_index], arr[i] = arr[i], arr[store_index]
                store_index += 1
        # drop pivot into sorted position
        arr[right], arr[store_index] = arr[store_index], arr[right]
        return store_index

    left, right = 0, len(arr) - 1
    while left < right:
        # random pivot avoids O(n^2) worst case
        pivot_index = random.randint(left, right)
        pivot_index = partition(left, right, pivot_index)
        if pivot_index == target_index:
            break
        elif pivot_index < target_index:
            left = pivot_index + 1
        else:
            right = pivot_index - 1
    return arr[target_index]

Key Functions & Tricks

• random.randint(left, right) — picks a pivot index, inclusive of both ends.
• Random pivot — avoids worst-case O(n^2) on sorted or adversarial input.
• list(scores) — shallow copy so in-place swaps don't mutate the caller's list.
• arr[i], arr[j] = arr[j], arr[i] — tuple-unpacking swap without a temp variable.
• Lomuto partition — swaps smaller-than-pivot elements left, then drops the pivot at store_index.
• target_index = len(arr) - k — kth largest is index n - k in ascending order.

How to Recognize This Pattern