NumPy按条件获取沿轴最大值索引 不满足阈值返回-1

原有代码问题 #

之前的实现存在逻辑漏洞：将所有≤0.9的元素替换为-1后，如果某一行没有元素大于0.9，整行值都会是-1，此时np.argmax遇到多个相等的最大值时，会默认返回第一个元素的索引0，因此输出了不符合预期的array([2, 0])。

正确简洁实现 #

不需要提前替换数组元素，先分别计算每行的最大值、每行最大值的索引，再通过阈值判断给不满足条件的行赋值-1即可：

import numpy as np
arr = np.array([[0.6,0.9,1],[0.3,0.5,0.7]])

result = np.where(arr.max(axis=1) > 0.9, arr.argmax(axis=1), -1)

运行上述代码可以得到预期输出array([ 2, -1])。

逻辑说明 #

• arr.max(axis=1)：沿列轴方向计算每一行的最大值，示例中得到结果[1. 0.7]
• arr.argmax(axis=1)：沿列轴方向取每一行最大值对应的列索引，示例中得到结果[2 0]
• np.where：逐行判断最大值是否大于0.9，满足条件则保留对应索引，不满足则直接赋值为-1，避免了全-1行取argmax得到错误索引的问题。

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