python - 使用 2D 数组进行 2D 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]])