python - 在numpy中，如何对一个与另一个顺序相同的数组进行排序？

Question

有两个numpy数组a和w，
两者都具有相同的形状 (d1,d2,..,dk,N)。
我们可以认为有 N 个形状为 (d1,d2,...,dk) 的样本。
现在，我想沿 a 的最后一个轴对 a 和 w 进行排序。
例如，a 和 w 的形状为 (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]])

我想要:

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

w 应该是:

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

当然，在这种情况下，下面的代码工作

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

但是当样本有很多维度(> 1)时，该代码不起作用。
任何人都可以提出解决方案吗？谢谢!

Answer 1

有一个函数，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]])
>>> 

当 a 时它也可以工作和 w具有任意形状。