Numpy 2D数组沿特定轴按其他2D索引

Question

我有一个二维数组：

>>> in_arr = np.array([[1,2],[4,3]])
array([[1, 2],
       [4, 3]])

我找到按列排序的索引，得到另一个2D数组：

>>> col_sort = np.argsort(in_arr, axis=1)
array([[0, 1],
       [1, 0]])

我想知道用第二个索引第一个的有效numpy切片：

>>> redordered_in_arr = np.*SOME_SLICE_METHOD*(in_arr, col_sort, axis=1)
array([[1, 2],
       [3, 4]])

这样做的目的是按列对数组执行一个（更复杂的）函数，例如：

>>> arr_with_function = reordered_in_arr ** np.array([1,2])
array([[1, 4],
       [3, 16]])

并将元素返回到其在数组中的原始位置

>>> return_order = np.argsort(col_sort, axis=1)
>>> redordered_in_arr = np.*SOME_SLICE_METHOD*(arr_with_function, return_order, axis=1)
array([[1, 4],
       [16, 3]]) 

好吧，在我打字的时候考虑一下，我可能会使用apply\u over\u axis，但是我仍然想知道如何有效地处理上面的问题，以防以后有价值。。你知道吗

Answer 1

如果您想就地执行所有这些操作，那么就不需要argsort()。Numpy在以下情况下支持就地操作：

In [12]: in_arr = np.array([[1,2],[4,3]])

In [13]: in_arr.sort(axis=1)

In [14]: in_arr **= [1, 2]

In [15]: in_arr
Out[15]: 
array([[ 1,  4],
       [ 3, 16]])

但是如果您需要排序项的索引，您可以通过简单的索引获得预期的结果。你知道吗

In [18]: in_arr[np.arange(2)[:,None], col_sort]
Out[18]: 
array([[1, 2],
       [3, 4]])