如何获取numpy条件筛选后数组argmax在原数组中的正确索引

当我们对numpy数组做条件筛选（比如a[a>5]）后，直接调用argmax()得到的是筛选后一维数组的索引，若直接用np.unravel_index映射回原数组形状，会得到错误结果——就像你示例里的(0, 2)，这并不是原数组中符合条件元素的最大值所在位置。下面是两种可靠的解决方法：

方法一：通过np.where获取筛选索引再定位 #

1. 先获取所有满足条件的元素在原数组中的索引：

   idx = np.where(a > 5)  # 返回两个数组，分别对应行、列索引
   

2. 找到筛选后数组中最大值的位置：

   filtered_max_idx = a[a>5].argmax()
   

3. 从筛选出的索引中取出对应原数组的坐标：

   original_max_idx = (idx[0][filtered_max_idx], idx[1][filtered_max_idx])
   

用你提供的数组测试：

import numpy as np
a = (np.random.random((10, 10))*10).astype(int)
# 假设a是你示例中的数组
idx = np.where(a > 5)
filtered_max_idx = a[a>5].argmax()
print(original_max_idx)  # 输出(0, 7)，正确对应原数组的最大值位置

方法二：掩码替换后直接取原数组argmax #

这种方法更简洁，通过将不满足条件的元素替换成极小值，让argmax自动忽略它们：

1. 创建掩码数组，把不符合条件的元素设为负无穷（确保比数组中所有元素都小）：

   masked_a = np.where(a > 5, a, -np.inf)
   

2. 直接对掩码数组取argmax并解析原数组索引：

   original_max_idx = np.unravel_index(masked_a.argmax(), a.shape)
   

测试代码：

masked_a = np.where(a > 5, a, -np.inf)
print(np.unravel_index(masked_a.argmax(), a.shape))  # 输出(0, 7)，正确

为什么原方法会出错 #

a[a>5]将符合条件的元素拉成了新的一维数组，它的argmax()返回的是这个新数组内的位置，而np.unravel_index需要的是原数组扁平化后的全局索引，两者的索引体系完全不同，因此直接映射会得到错误结果。

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