In numpy, how to sort an array with the same order with another one?

user_13859092 0
Viewed 130

There are two numpy array a and w, both of which have the same shape (d1,d2,..,dk,N). We can think there are N sample with shape (d1,d2,...,dk).

Now, I want to sort a and w along a's last axis.

For example, a and w have shape (2,4):

a = [[3,2,4,1],
     [2,3,1,4]]
w = [[10,20,30,40],
     [80,70,60,50]]
sorted_index = a.argsort()
# array([[3, 1, 0, 2],
#        [2, 0, 1, 3]])

I want:

a = a.sort() # default axis = -1
# a = [[1,2,3,4],
#      [1,2,3,4]]

and w should be:

# w = [[40,20,10,30],
#      [60,80,70,50]]

Of course, in that case, the following code work

x = a.argsort()
w[0,:] = w[0,x[0]]
w[1,:] = w[1,x[1]]

But when the sample have many dimension (>1), that code doesn't work. Can anyone come up with solutions? Thanks!

numpy sorting python
1 Answers

There's a function for that, np.take_along_axis:

>>> a = np.array([[3,2,4,1], [2,3,1,4]])
>>> w = np.array([[10,20,30,40], [80,70,60,50]])
>>> sorted_index = a.argsort()
>>> sorted_index
array([[3, 1, 0, 2],
       [2, 0, 1, 3]])
>>> np.take_along_axis(a, sorted_index, axis=-1)
array([[1, 2, 3, 4],
       [1, 2, 3, 4]])
>>> np.take_along_axis(w, sorted_index, axis=-1)
array([[40, 20, 10, 30],
       [60, 80, 70, 50]])
>>>

It will also work when a and w have arbitrary shape.

user_5217760
Related