如何通过numpy argmax返回的索引获取最大值？含N维数组方案

这问题我之前踩过坑，直接用u[e]确实得不到想要的结果，因为numpy的索引规则和我们直觉里的不一样。下面分两种情况给你讲清楚怎么解决：

一、3维数组的具体解决方案 #

先拿你的示例代码来说：

import numpy as np
u = np.arange(12).reshape(3,4,1)
e = u.argmax(axis=2)

你得到的e是形状为(3,4)的数组，但原数组u是(3,4,1)，直接用u[e]会把e里的每个元素当作第一个轴的索引，所以会重复取u[0]的内容，完全不是我们要的沿axis=2的最大值。

方法1：用np.take_along_axis（最推荐） #

这个函数就是专门用来配合argmax/argmin获取对应极值的，用法很直接：

# 给索引数组e增加一个维度，对应axis=2的位置，让它的形状和原数组匹配
expanded_e = e[..., np.newaxis]  # 或者用e[..., None]，效果一样
max_vals = np.take_along_axis(u, expanded_e, axis=2)
print(max_vals)
# 输出：
# [[[ 0]
#   [ 1]
#   [ 2]
#   [ 3]]
#
#  [[ 4]
#   [ 5]
#   [ 6]
#   [ 7]]
#
#  [[ 8]
#   [ 9]
#   [10]
#   [11]]]

如果想去掉长度为1的维度，可以再加个squeeze：

max_vals_squeezed = max_vals.squeeze(axis=2)
print(max_vals_squeezed)
# 输出：
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

方法2：手动构造多轴索引 #

如果你想理解底层逻辑，可以用np.indices生成其他轴的索引，然后组合起来：

# 生成axis=0和axis=1的索引数组，形状和e一致
i, j = np.indices(e.shape)
# 用(i,j,e[i,j])来索引原数组u
max_vals = u[i, j, e]
print(max_vals)
# 输出和上面squeeze后的结果一样：
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

二、N维数组的通用解法 #

不管你的数组是几维，np.take_along_axis都能通用，核心就是给argmax得到的索引数组在指定axis上增加一个维度，然后传入函数即可。这里给你封装成一个通用函数：

def get_max_from_argmax(arr, axis):
    argmax_indices = arr.argmax(axis=axis)
    # 扩展索引数组的维度，匹配原数组的轴
    expanded_indices = np.expand_dims(argmax_indices, axis=axis)
    # 获取对应最大值
    max_values = np.take_along_axis(arr, expanded_indices, axis=axis)
    # 可选：去掉长度为1的维度，根据需求决定是否保留
    # max_values = max_values.squeeze(axis=axis)
    return max_values

举个4维数组的例子测试一下：

# 生成一个4维随机数组
arr_4d = np.random.rand(2, 3, 4, 5)
# 沿axis=2取最大值
max_vals_4d = get_max_from_argmax(arr_4d, axis=2)
print(f"原数组形状：{arr_4d.shape}")
print(f"最大值数组形状：{max_vals_4d.shape}")
# 输出：
# 原数组形状：(2, 3, 4, 5)
# 最大值数组形状：(2, 3, 1, 5)

如果你不需要那个长度为1的维度，打开函数里的squeeze注释即可，得到的形状就是(2,3,5)。

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