如何在NumPy数组的最后一个轴上进行元素选取？

方法1：高级索引（最推荐） #

通过构造前两个维度的网格索引，结合idx实现精准取值，这是最直观且性能最优的方式：

import numpy as np

# 示例数据
shape = (4,3,2)
x = np.random.uniform(0,1, shape)
idx = np.random.randint(0,shape[-1], shape[:-1])

# 生成前两维的网格索引
i, j = np.indices(shape[:-1])
# 高级索引取值
result = x[i, j, idx]

如果想节省内存（尤其处理大数组时），可以用np.ogrid生成开放式网格索引，无需创建完整的网格数组：

i, j = np.ogrid[:shape[0], :shape[1]]
result = x[i, j, idx]

方法2：np.take_along_axis（NumPy 1.15+ 支持） #

这个API专门用于沿指定轴按索引提取元素，语法更简洁：

# 给idx增加最后一维，匹配原数组维度
expanded_idx = idx[..., np.newaxis]
# 沿最后一维提取后，去除多余维度
result = np.take_along_axis(x, expanded_idx, axis=-1).squeeze(axis=-1)

方法3：扁平化索引转换 #

通过多维索引转一维扁平索引的方式取值，适合理解底层索引逻辑：

i, j = np.indices(shape[:-1])
# 计算每个位置的一维扁平索引
flat_idx = np.ravel_multi_index((i, j, idx), shape)
# 取值后恢复原维度
result = x.flat[flat_idx].reshape(shape[:-1])

以上所有向量化方法都基于NumPy的C级优化实现，完全避免了Python循环的性能开销，在数组规模较大时，性能提升非常显著。

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