用于多维ndarray的argsort

Question

我正在尝试获取索引以按最后一个轴对多维数组进行排序，例如

>>> a = np.array([[3,1,2],[8,9,2]])

我想要这样的索引i，

>>> a[i]
array([[1, 2, 3],
       [2, 8, 9]])

根据numpy.argsort的文档，我认为应该执行此操作，但出现错误：

>>> a[np.argsort(a)]
IndexError: index 2 is out of bounds for axis 0 with size 2

编辑：我需要重新排列相同形状的其它阵列（例如阵列b，使得a.shape == b.shape以相同的方式）…以便

>>> b = np.array([[0,5,4],[3,9,1]])
>>> b[i]
array([[5,4,0],
       [9,3,1]])

Answer 1

上面的答案现在已经过时了，因为在numpy
1.15中添加了新功能以使其更简单。take_along_axis（https://docs.scipy.org/doc/numpy-1.15.1/reference/generation/numpy.take_along_axis.html）允许您执行以下操作：

>>> a = np.array([[3,1,2],[8,9,2]])
>>> np.take_along_axis(a, a.argsort(axis=-1), axis=-1)
array([[1 2 3]
       [2 8 9]])