Used in **LLM token sampling** — before sampling the next token, we restrict the vocabulary to only the top-K most likely tokens.
Given a list of logits, return the **indices** of the top-K highest values, ordered from highest to lowest.
top_k_indices([5, 3, 1, 4, 2], k=3) → [0, 3, 1] # Values: 5=idx0, 4=idx3, 3=idx1
Similar Problems
Test Cases (2 visible · 2 hidden)
Case 1: Basic top-3
Input: top_k_indices([5,3,1,4,2], 3)
Expected: [0, 3, 1]
Case 2: Ascending input
Input: top_k_indices([1,2,3,4,5], 2)
Expected: [4, 3]
⌘↵ Run · ⌘⇧↵ Submit