Numpy中整数数组索引的广播与对齐高效实现问询

假设有形状为(n, d)的Numpy数组a，示例代码如下：

np.random.seed(1)

n, d = 5, 3
a = np.random.randn(n, d)

另有形状为(m, n)的整数索引数组indices，元素取值范围为0到d-1，用于索引a的第二维度，示例代码如下：

m = 10
indices = np.random.randint(low=0, high=d, size=(m, n))

需求是用indices对a的第二维度做索引，实现按n维度对齐、按m维度批量处理。现有两种解决方案但写法不够简洁，希望找到更优雅的Numpy风格广播实现方式（类似a.T[indices]的写法），且要求在m和n取值较大时性能最优。

方案1：通过列表推导循环拼接，再转置得到结果

result = np.vstack([a[i, :][indices[:, i]] for i in range(n)]).T
print(result.shape)
# (10, 5)

方案2：利用转置和对角线提取

np.diagonal(a.T[indices], axis1=1, axis2=2)

推荐使用np.take_along_axis函数，这是Numpy专门为这类轴对齐索引场景设计的API，完全基于向量化操作，性能最优且可读性强：

# 将a扩展为(1, n, d)，indices扩展为(m, n, 1)，沿第2轴取对应元素后压缩维度
result = np.take_along_axis(a[np.newaxis, :, :], indices[:, :, np.newaxis], axis=2).squeeze(axis=2)

或者用更简洁的维度匹配索引写法，同样是纯向量化操作：

# 利用np.arange(n)匹配a的第一维度，结合indices的转置实现对齐索引
result = a[np.arange(n), indices.T].T

两种写法都避免了Python层面的循环，完全依赖Numpy的底层优化，在m和n较大时性能远优于方案1，和方案2性能相当但可读性更强。

内容的提问来源于stack exchange，提问作者Leon