使用2D数组进行2D Numpy Array索引

Question

  position_array = np.array([[1,2]
                      [0,1]
                      [0,2]])

  original_array = np.array([[5,6,7]
                             [7,8,2]
                             [10,6,2]])

result_i_want = np.array([[6,7]
                        [7,2]
                        [10,2]])

该怎么做？

从original_array获取一个新数组，并将position_array用作索引??

例如，当position_array的第一行为[1,2]时，意味着从original_array的第一行中获取列1和列2

如何以最简单的方式进行操作？

Answer 1

您要np.take_along_axis，设置axis=1：

np.take_along_axis(original_array, position_array, axis=1)

array([[ 6,  7],
       [ 7,  8],
       [10,  2]])

---

或者您也可以使用advanced indexing：

original_array[np.arange(original_array.shape[0])[:,None], position_array]

array([[ 6,  7],
       [ 7,  8],
       [10,  2]])