如何在Numpy中获取大于阈值的前N个最大值的索引？

直接使用全向量化操作替代 Python 层循环，性能远高于原生循环写法，尤其适合大尺寸矩阵场景：

import numpy as np

# 取每行前N大的元素索引，argpartition无需全排序，效率高于argsort
top_n_idx = np.argpartition(results, -N, axis=1)[:, -N:]
# 生成符合阈值要求的掩码
mask = np.take_along_axis(results, top_n_idx, axis=1) > threshold

# 【方案1：扁平化输出】直接过滤得到所有符合要求的索引
flat_indexes = top_n_idx[mask]

# 【方案2：按行分组输出】和示例输出格式一致
row_indexes = [row[mask[i]] for i, row in enumerate(top_n_idx)]

实现说明 #

• 用np.argpartition替代np.argsort，不需要对整行元素完全排序，仅需定位到前N大的元素即可，每行时间复杂度从O(m log m)降到O(m)，矩阵列数越大性能优势越明显
• 所有数值计算、比较逻辑都在Numpy底层完成，避免了Python层循环的性能损耗，大矩阵下性能比原始写法提升数十到上百倍

示例验证 #

用你给出的测试用例运行：

results = np.array([
    [3, 2, 6, 7, 0],
    [4, 1, 6, 4, 0],
    [7, 10, 6, 9, 8]
])
N = 3
threshold = 5

得到的row_indexes输出为：

[array([2, 3]), array([2]), array([1, 3, 4])]

完全符合预期要求，且索引顺序可根据需要自行排序调整。

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