如何利用numpy.argmax得到的索引批量获取二维数组对应值？

要实现你的需求，核心是结合行索引数组和argmax得到的列索引数组来定位元素，而不是直接用列索引切片。以下是具体实现：

示例代码 #

import numpy as np

# 构造匹配预期输出的示例数组
a = np.array([[1, 3, 5, 7], [9, 2, 4, 6], [3, 8, 2, 5], [6, 1, 2, 7]])
b = np.array([[2, 4, 6, 8], [5, 3, 1, 7], [4, 8, 6, 2], [3, 5, 1, 7]])

# 获取a每行最大值的列索引
row_max_indices = np.argmax(a, axis=1)

# 生成行索引数组，批量提取对应元素
result = b[np.arange(a.shape[0]), row_max_indices]
print(result)  # 输出：array([8, 5, 8, 7])

关键说明 #

• np.arange(a.shape[0])生成与数组行数一致的行索引序列（比如4行数组会生成[0,1,2,3]），和row_max_indices（每行最大值的列索引数组）组合后，Numpy会按(行0,列idx0), (行1,列idx1)...的规则逐个定位元素，直接得到一维结果数组。
• 你之前用b[:, row_max_indices]的错误在于：这种写法是取所有行的row_max_indices指定列，得到的是二维数组，而非每行对应单个元素的一维结果，不符合需求。

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