如何沿指定轴按最大值位置切片多维NumPy数组

完全可以用NumPy原生的高级索引机制替代显式循环，写法简洁且性能远高于Python层循环，是标准的numpythonic实现。

核心逻辑 #

你原有循环的本质是对第二个轴（长度为m）的每个位置i，配对取第一个轴上wheremax[i]位置的全部第三轴数据。利用NumPy整数数组的高级索引规则：当传入多个形状相同的整数数组作为索引时，会按位置一一对应提取元素，正好匹配这个需求。

实现代码 #

# 保留你原有的最大值索引计算逻辑
sub = a[:, :, j]
wheremax = np.argmax(sub, axis=0)  # 形状为(m,)

# 一行代码替换整个for循环
new_arr = a[wheremax, np.arange(m), :]

代码里的索引对应关系：

• 第一个位置传wheremax：指定每个配对位置要取的第一轴索引
• 第二个位置传np.arange(m)：生成0到m-1的序列，和wheremax一一对应，指定第二轴的索引
• 第三个位置传:：保留第三轴的所有元素

最终new_arr的形状正好是(m, l)，和你循环实现的结果完全一致。

可选紧凑写法 #

如果要进一步精简代码，可以省略单独提取sub的步骤，直接在原数组上计算最大值索引：

wheremax = np.argmax(a[..., j], axis=0)
new_arr = a[wheremax, np.arange(a.shape[1]), :]

这里...是NumPy的省略号语法，自动匹配前面所有维度，a[..., j]和你写的a[:, :, j]效果完全相同。

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