如何从包含最多3个值的2d numpy数组中获取列的索引

Question

I have a array:

a = np.array([[22,11,44,33,66],
              [22,11,2,1,66],
              [1,11,44,22,4],
              [22,11,88,99,66]])

and as a output I want a array containing index of max 3 values as 2d array. For example for above array output would be:

array([[4,2,3],
       [4,0,1],
       [2,3,1],
       [3,2,4]])

Answer 1

To get the top k elements of an array, partition it. Since partitioning normally gives you the k lowest elements, use the reverse indices:

k = 3
top = np.argpartition(a, -k, axis=1)[:, -k:]

If you need to have the indices sorted in descending order, use np.argsort with the result:

rows = np.arange(a.shape[0])[:, None]
s = np.argsort(a[rows, top], axis=1)[:, ::-1]
top = top[rows, s]

rows is necessary to make sure that all the indices are selected properly when you do fancy indexing with top and s . The indices for each row have to be reversed ( [:, ::-1] ) to get ascending order.