Indexes of top N values in 2 dimensional array

Question

I have a following 2 dimensional numpy array:

[[2,3,1,5], [4,3,6,5], [7,10,8,9]]

For each array on axis=1 I'd like to get indexes of N items with highest values. Example: response for N=2 should look like this:

[[3,1], [2,3], [1,3]]

Possible duplicate of [N largest values in each row of ndarray](https://stackoverflow.com/questions/30332908/n-largest-values-in-each-row-of-ndarray) — Mr. T, Feb 06 '18 at 13:51

score 1 · Accepted Answer · answered Feb 06 '18 at 13:18

1

a=[[2,3,1,5], [4,3,6,5], [7,10,8,9]]


def fetchindex(array,N):
    result=[]
    for row in array:
        enumeratedrow=enumerate(row)
        sortedarray=sorted(enumeratedrow,key=lambda x:x[1],reverse=True)[:N]
        temp=[i[0] for i in sortedarray]
        result.append(temp)
    return result 

print fetchindex(a,2)

RESULT

[[3, 1], [2, 3], [1, 3]]

answered Feb 06 '18 at 13:18

Albin Paul

3,330
2
14
30

Thanks but I was hoping for much cleaner/shorter solution, e.g. using numpy.argsort – filip Feb 06 '18 at 13:45

Indexes of top N values in 2 dimensional array

1 Answers1