Python 使用numpy数组索引表示列表理解

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

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

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

result = P[MAGIC_HERE]

---

如果是这样的话，哪一种计算成本更低（更快）？

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

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

---

那相当快。谢谢