2

I got a numpy 2D array, and the list of indices corresponding to the top 3 elements obtained using argsort. Now, I am trying to extract the values corresponding to this indices, and it is not working. What is the workaround ?.

A = array([[0.19334242, 0.9787497 , 0.41453434, 0.35298119, 0.17943745,
        0.63468207, 0.43840688],
       [0.39811914, 0.68040634, 0.7589702 , 0.3573046 , 0.16365397,
        0.86329535, 0.48559053],
       [0.5848541 , 0.54203383, 0.27262654, 0.21979374, 0.06917679,
        0.10586995, 0.57083441],
       [0.76765549, 0.05703751, 0.83383973, 0.71867625, 0.16338699,
        0.85721418, 0.5953548 ]])

np.flip(A.argsort(),axis=1)[:,0:3]
array([[1, 5, 6],
       [5, 2, 1],
       [0, 6, 1],
       [5, 2, 0]])

gets error

>>> A[np.flip(A.argsort(),axis=1)[:,0:3]]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: index 5 is out of bounds for axis 0 with size 4
Shew
  • 1,557
  • 1
  • 21
  • 36
  • @hpaulj Yes, I looked at argsort. I did not get the comment in full sense ?. Is there a way, I can get the values directly using argosrt in case of 2D arrays ? – Shew Mar 01 '18 at 08:15
  • Top 3 element in what dimension? rows, columns, flattened `A`? – hpaulj Mar 01 '18 at 08:17
  • `A[idx]` selects rows, not elements from rows. – hpaulj Mar 01 '18 at 08:18
  • The way to apply the index is `out = a[np.arange(a.shape[0])[:,None],idx]`. This is from the second answer of the duplicate. – hpaulj Mar 01 '18 at 12:12

1 Answers1

0
In [22]: A.ravel()[A.argsort(axis=None)[::-1][:3]]
Out[22]: array([ 0.9787497 ,  0.86329535,  0.85721418])

Explanation

By default, argsort() sorts along the last axis. In your case, you want to sort a flattened version of the array as you don't give any meaning to the fact that the array is 2D. This happens by passing axis=None to argsort().

Since you get 1D indices, you also need to access values on a flattened version of the array, which is what ravel() do.

[::-1] reverses the argsort array to get top values first and [:3] gets the first 3 values.

Note: there are other and possibly more efficient ways to do that, but this was the first thing that came to my mind.

nicoco
  • 1,421
  • 9
  • 30