1

I have a numpy 2d array of real numbers, for example

A=
np.array(
[[0.1, 0.01, 0.4, 0.9], 
[0.0005, -0.2, -0.1, 0.6], 
[-0.3, -0.5, 0.2, 0.9]])

and a vector of indices of the same size of A.shape[1]: idx=[5, 2, 3, 9]

for every row of A, I need to sort the entries from high to low, and provide the corresponding elements from idx. For example, in the case above the answer should be:

np.array([[9, 3, 5, 2], [9, 5, 3, 2], [9, 3, 5, 2]])

There can be millions of row in A. What is the most efficient way of doing this?

Bob
  • 10,741
  • 27
  • 89
  • 143

1 Answers1

3

You can use np.argsort() to get the indices of your array in a sorted mode and reverse it to get the indices of decreasing mode, then use a simple indexing to get your expected output:

>>> A.argsort()
array([[1, 0, 2, 3],
       [1, 2, 0, 3],
       [1, 0, 2, 3]])
>>> idx=np.array([5, 2, 3, 9])

>>> indices = A.argsort()[:,::-1]
>>> 
>>> idx[indices]
array([[9, 3, 5, 2],
       [9, 5, 3, 2],
       [9, 3, 5, 2]])
Mazdak
  • 105,000
  • 18
  • 159
  • 188