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

Question

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

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

是否可以使用数组索引来表达这一点？也就是说，是否存在形式的等效表达式：

result = P[MAGIC_HERE]

如果是这样，那么计算成本会更低（更快）？

Answer 1

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

data