python - numpy.argmin用于大于阈值的元素

我感兴趣的是得到满足特定条件（在我的例子中，是中等阈值）的一维NumPy数组中最小值的位置。例如：

import numpy as np

limit = 3
a = np.array([1, 2, 4, 5, 2, 5, 3, 6, 7, 9, 10])

我想有效地屏蔽a中所有低于限制的数字，这样np.argmin的结果将是6。是否有一种计算成本低廉的方法来屏蔽不满足条件的值，然后应用np.argmin？

最佳答案

您可以存储有效的索引，并使用那些既从“cc>选择有效元素，也可以在选定元素中索引为a的索引，以得到最终的索引输出。因此，实现将如下所示-

valid_idx = np.where(a >= limit)[0]
out = valid_idx[a[valid_idx].argmin()]

样本运行-

In [32]: limit = 3
    ...: a = np.array([1, 2, 4, 5, 2, 5, 3, 6, 7, 9, 10])
    ...:

In [33]: valid_idx = np.where(a >= limit)[0]

In [34]: valid_idx[a[valid_idx].argmin()]
Out[34]: 6

运行时测试-
对于性能基准测试，在本节中，我将argmin()与本文前面针对各种数据大小提出的基于常规阵列的解决方案进行比较。

def masked_argmin(a,limit): # Defining func for regular array based soln
    valid_idx = np.where(a >= limit)[0]
    return valid_idx[a[valid_idx].argmin()]

In [52]: # Inputs
    ...: a = np.random.randint(0,1000,(10000))
    ...: limit = 500
    ...:

In [53]: %timeit np.argmin(np.ma.MaskedArray(a, a<limit))
1000 loops, best of 3: 233 µs per loop

In [54]: %timeit masked_argmin(a,limit)
10000 loops, best of 3: 101 µs per loop

In [55]: # Inputs
    ...: a = np.random.randint(0,1000,(100000))
    ...: limit = 500
    ...:

In [56]: %timeit np.argmin(np.ma.MaskedArray(a, a<limit))
1000 loops, best of 3: 1.73 ms per loop

In [57]: %timeit masked_argmin(a,limit)
1000 loops, best of 3: 1.03 ms per loop