无循环提取三维NumPy数组最后轴最大值（兼容CuPy）

针对你的需求（无循环、兼容CuPy语法），最直接且高效的方法是使用np.take_along_axis（CuPy中对应cp.take_along_axis，语法完全一致），它专门用于根据沿指定轴的索引提取对应元素，完美适配你的场景。

实现代码 #

import numpy as np

arr = np.array(
    [[[ 6, -2, -6, -5],
      [ 1, 12,  3,  9],
      [21,  7,  9,  8]],

     [[15, 12, 20, 12],
      [17, 15, 17, 23],
      [22, 18, 27, 32]]])

indexes = arr.argmax(axis=2, keepdims=True)
result = np.take_along_axis(arr, indexes, axis=2)
print(result)

输出结果 #

array([[[ 6],
        [12],
        [21]],

       [[20],
        [23],
        [32]]])

原理说明 #

• take_along_axis会沿着指定的axis=2（最后一轴），用indexes中的索引值提取原数组arr对应位置的元素
• 因为你在argmax时设置了keepdims=True，保证了indexes的维度和arr一致，所以提取后的结果维度也会和预期匹配
• 该方法完全基于向量操作，没有循环，能充分利用GPU并行性（CuPy环境下直接替换np为cp即可运行）

补充：手动构造索引的方法（可选） #

如果你想理解底层索引逻辑，也可以手动构造前两个轴的索引，结合广播实现：

# 构造前两个轴的索引
i, j = np.indices(arr.shape[:2])
# 组合索引并提取元素
result = arr[i, j, indexes.squeeze()]
# 重新恢复最后一维
result = result[..., np.newaxis]

不过这种方法不如take_along_axis简洁，优先推荐前者。

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