如何在特定轴上索引由索引给定的多维数组？

问题描述

假设我有一个形状为A的4d数组(D0,D1,D2,D3)。我有一个形状为B的一维数组(D0,)，其中包含我在第2轴处需要的索引。

实现我所需的简单方法：

output_lis = []
for i in range(D0):
    output_lis.append(A[i,:,B[i],:])
#output = np.concatenate(output_lis,axis=0) #it is wrong to use concatenate. Thanks to @Mad Physicist. Instead,using stack.
output = np.stack(output_lis,axis=0) #shape: [D0,D3]

所以，我的问题是如何用numpy API快速实现它？

解决方法

使用花式索引在锁步中沿二维移动。在这种情况下，arange提供序列i，而B提供序列B[i]：

A[np.arange(D0),:,B,:]

与(D0,D1,D3)循环结果的形状不同，此数组的形状确实为for。

要从示例中获得相同的结果，请使用stack（添加新轴），而不要使用concatenate（使用现有轴）：

output = np.stack(output_lis,axis=0)