2

I have a numpy array a like this

In [318]: a
Out[318]:
array([[0. , 1. , 2. , 3. ],
       [0.5, 0.3, 0.2, 0.25]])

I need to sort along the second row (the one with [0.5,0.3,0.2,0.25]), while having the first row changed accordingly. In this case, the expected result is

2    3     1.   0
0.2, 0.25, 0.3, 0.5

How can I do this? Thank you. I tried np.sort with axis=-1 and 0; they are not what I need.

Important note: The performance is the key in my problem solving. My arrays, from an image processing application, are generally of N columns with an N close to 4 million.

zell
  • 9,830
  • 10
  • 62
  • 115

1 Answers1

2

Use np.argsort() to get the index of sorted row and then use it as a mask to sort the whole array (column based):

In [69]: mask = np.argsort(a[1])                                                                                                                                                                            

In [70]: a[:, mask]                                                                                                                                                                                         
Out[70]: 
array([[2.  , 3.  , 1.  , 0.  ],
       [0.2 , 0.25, 0.3 , 0.5 ]])
Mazdak
  • 105,000
  • 18
  • 159
  • 188