使用 numpy 数组作为另一个数组的第二个 dim 的索引？[重复]

问题描述：

例如，我有两个 numpy 数组，

A = np.array(
  [[0,1], 
   [2,3], 
   [4,5]])
B = np.array(
  [[1],
   [0],
   [1]], dtype='int')

并且我想从的每一行中提取一个元素A，并且该元素由索引B，所以我想要以下结果：

C = np.array(
  [[1],
   [2],
   [5]])

我试过了A[:, B.ravel()]，但它会广播B，这不是我想要的。还研究了一下np.take，似乎不是解决我问题的正确方法。

不过，我可以np.choose通过转置来使用A，

np.choose(B.ravel(), A.T)

但还有其他更好的解决方案吗？

解决方案 1：

您可以使用NumPy's purely integer array indexing-

A[np.arange(A.shape[0]),B.ravel()]

样本运行 -

In [57]: A
Out[57]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

In [58]: B
Out[58]: 
array([[1],
       [0],
       [1]])

In [59]: A[np.arange(A.shape[0]),B.ravel()]
Out[59]: array([1, 2, 5])

请注意，如果B是1D此类列索引的数组或列表，则可以简单地跳过展平操作.ravel()。

样本运行 -

In [186]: A
Out[186]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

In [187]: B
Out[187]: [1, 0, 1]

In [188]: A[np.arange(A.shape[0]),B]
Out[188]: array([1, 2, 5])

解决方案 2：

C = np.array([A[i][j] for i,j in enumerate(B)])