Индексирование массива двумерных массивов другими 2D вдоль определенной оси

Question

У меня есть 2D-массив:

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

и я нахожу отсортированные индексы по столбцам, чтобы получить другой двумерный массив:

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

Я хотел бы знать эффективный кусочек для индексации первого по второму:

>>> 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_over_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]])