使用numpy数组索引的快速列表理解

Question

我有一个三维数组P，形状（32,2,10） - 32个样本，每个形状（2,10） - 和一个1d数组A，形状（32 ）。 A中的值为0或1，用于选择每个样本的第一行或第二行。我使用列表理解工作的解决方案，如下图所示：

result = np.array([P[i, A[i]] for i in range(32)]) 

是否有可能表达这种只使用数组索引？也就是说，有如下形式的等效表达式：

result = P[MAGIC_HERE] 

如果是这样，这将是较少计算上是昂贵的（更快）？

Answer 1

您可以使用numpy的的索引像这样：

result = P[np.arange(32), A[np.arange(32)]]