使用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)]]