Top-K Values per Row

Problem

Return the k largest values of each row of A, sorted in descending order.

If k exceeds the row length, return all values. Result entry i is the top-k list for row i.

Hint

Sort each row descending, then slice the first k entries.

case 1

in: [[[3,1,2],[6,5,4]],2]

expected: [[3,2],[6,5]]

case 2

in: [[[1,2,3]],3]

expected: [[3,2,1]]

case 3

in: [[[5,5,4]],2]

expected: [[5,5]]

case 4

in: [[[1,2],[3,4]],1]

expected: [[2],[4]]

Starter codePython

def top_k_rows(A, k):
    # Your code here
    pass

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