使用二维数组进行二维 Numpy 数组索引 [重复]

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