使用 numpy 数组作为另一个数组的第二个暗淡的索引？ [复制]答案

【问题标题】：using an numpy array as indices of the 2nd dim of another array? [duplicate]使用 numpy 数组作为另一个数组的第二个暗淡的索引？ [复制]
【发布时间】：2016-06-18 13:54:28
【问题描述】：

例如，我有两个 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，似乎不是解决我的问题的正确方法。

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

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

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

【问题讨论】：

可能重复：stackoverflow.com/q/37878946/190597
@unutbu 这里有点不同，因为我们在每行选择一个元素，而不是在链接的上一个问题中每行选择多个元素。
@unutbu，你链接的答案使用np.take，我认为它不能解决我的问题？

标签： python numpy indexing

【解决方案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])

【讨论】：

所以根本没有广播？
好吧，如果我从问题中正确得到了预期的输出，那么您不需要在这里广播。我们基本上在每个暗处使用整数来选择元素。我们有来自B 的第二个暗淡索引，所以我们只需要使用np.arange 为第一个暗淡创建相应的索引。希望这是有道理的！

【解决方案2】：

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

【讨论】：