How to get index of columns from 2d numpy array containing max 3 values

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]])

@Josh. It's pretty clear, unless it's me missing something: each output row contains the indices of the top 4 elements of the corresponding input row. — Mad Physicist, Nov 11 '18 at 17:32
@MadPhysicist Yes, indices of top 3 elements. Your solution below is not giving desired result — saurabh kumar, Nov 11 '18 at 18:07
[How do I get indices of N maximum values in a NumPy array?](https://stackoverflow.com/questions/6910641/how-do-i-get-indices-of-n-maximum-values-in-a-numpy-array) — Nirmal, Nov 12 '18 at 06:40

Mad Physicist · Accepted Answer · 2018-11-12T06:03:06.313

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.

How to get index of columns from 2d numpy array containing max 3 values

1 Answers1