MediumNLPnlp-382

Top K Token Indices

Problem

Return the indices of the k highest scores, sorted by score descending with ties broken by index ascending.

Given the scores and k, return the index list (at most k entries).

Hint

Sort with a tuple key that negates the score for descending order.

case 1

in: [[0.1,0.9,0.5],2]

expected: [1,2]

case 2

in: [[1,0.5,1],2]

expected: [0,2]

case 3

in: [[0.2],3]

expected: [0]

Starter codePython

def top_k_token_indices(scores, k):
    # Your code here
    pass

3 test cases · run your code against them in the editor