如何在NumPy中根据每行最大值获取对应的列索引？

嘿，这个需求其实用NumPy内置的函数就能轻松解决！咱们先看你的具体场景：

你有这样一个NumPy数组：

import numpy as np
myArray = np.array([[ 3, 8, 0, 0], [ 9, 7, 8, 1], [ 3, 8, 9, 7]])

想要获取每行最大值对应的列索引，预期结果是[1, 0, 2]。

最简单的解决方案：使用np.argmax() #

NumPy的argmax()函数专门用来返回数组中最大值的索引，我们只需要指定axis=1参数，告诉它**沿着每行（横向）**去查找：

# 获取每行最大值的列索引
max_indices = np.argmax(myArray, axis=1)
# 转成列表格式（如果需要的话）
result = max_indices.tolist()
print(result)  # 输出: [1, 0, 2]

补充说明 #

• axis=1是关键：如果不指定axis，argmax()会默认把数组扁平化，返回整个数组最大值的索引，这显然不是我们要的。指定axis=1后，函数会逐行处理，返回每行最大值的列位置。
• 如果某一行存在多个相同的最大值，argmax()只会返回第一个出现的最大值的列索引哦。比如如果有一行是[5, 9, 9]，它会返回索引1而不是2。

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