Top-K Scores

Return indices of the k largest scores in descending score order.

Start with the concrete trace below. It shows the state the algorithm must carry as it runs.

Problem trace

Top-K Scores: Use partial partition only to select top-k membership, then sort those candidate indices by their scores for the required descending order.

Input and goalReturn indices of the `k` largest scores in descending score order.

Validate k against score countFor scores [0.1,0.9,0.4,0.8] and k=2, the guard 1<=2<=4 passes.

input start0.100.910.420.83

guard1 <= k=2 <= len(scores)=4

Recognize it

Use argpartition when k ranked elements are needed from a much larger score vector and fully sorting all n scores is unnecessary.

Keep true

After partition, every selected candidate belongs to the top-k score group although candidate order is unspecified; after subset argsort reversal, those same indices are descending by score.

Reuse it

Separate selection from ordering: first isolate the small set that can contain the answer, then pay sorting cost only on that set.

Pattern: Partial selection, then sort only the selected values.

Simple idea: argpartition finds the top group without sorting every score. Sort the small selected group for final order.

import numpy as np

def top_k_indices(scores: np.ndarray, k: int) -> np.ndarray:
   if not 1 <= k <= len(scores):
      raise ValueError("k must name an item in scores")
   candidates = np.argpartition(scores, -k)[-k:]
   return candidates[np.argsort(scores[candidates])[::-1]]

Cost: $O(n+k\log k)$ average time and $O(k)$ selected space.