如何利用数组的最大值索引从另一数组取值？尝试代码未得预期结果

已知数组：

import numpy as np
a = np.array([[11,12,39,14], [15,16,29,18]])
b = np.array([[1,2,3,4], [5,6,7,8]])

你的代码result = b[range(len(a)), a.argmax(0)]出错的原因是：

• range(len(a))生成的是长度为2的行索引数组[0,1]，而a.argmax(0)生成的是长度为4的列索引数组[1,1,0,1]，两者维度不匹配，索引逻辑不符合需求。

正确的做法是，针对每一列，用a.argmax(0)获取该列最大值所在的行索引，再用np.arange(a.shape[1])获取列索引，以此从b中提取对应元素：

import numpy as np
a = np.array([[11,12,39,14], [15,16,29,18]])
b = np.array([[1,2,3,4], [5,6,7,8]])

# 获取每列最大值的行索引
row_indices = a.argmax(axis=0)
# 获取列索引数组
col_indices = np.arange(a.shape[1])
# 提取对应元素
result = b[row_indices, col_indices]
print(result)  # 输出: [5 6 3 8]

也可以简化为一行代码：

result = b[a.argmax(0), np.arange(a.shape[1])]

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