python - 如何仅替换 numpy 数组中大于某个值的前 n 个元素？

Question

我有一个这样的数组 myA:

array([ 7,  4,  5,  8,  3, 10])

如果我想用 0 替换所有大于值 val 的值，我可以简单地这样做:

myA[myA > val] = 0

这给了我想要的输出(对于 val = 5):

 array([0, 4, 5, 0, 3, 0])

但是，我的目标不是全部替换，而是仅替换此数组中大于值 val 的前 n 个元素。

因此，如果 n = 2，我想要的结果将如下所示(10 是第三个元素，因此不应被替换):

array([ 0,  4,  5,  0,  3, 10])

一个简单的实现是:

import numpy as np

myA = np.array([7, 4, 5, 8, 3, 10])
n = 2
val = 5

# track the number of replacements
repl = 0

for ind, vali in enumerate(myA):

    if vali > val:

        myA[ind] = 0
        repl += 1

        if repl == n:
            break

这行得通，但也许有人可以想出一种巧妙的掩蔽方式!？

Answer 1

以下应该有效:

myA[(myA > val).nonzero()[0][:2]] = 0

自 nonzero将返回 bool 数组 myA > val 非零的索引，例如正确。

例如:

In [1]: myA = array([ 7,  4,  5,  8,  3, 10])

In [2]: myA[(myA > 5).nonzero()[0][:2]] = 0

In [3]: myA
Out[3]: array([ 0,  4,  5,  0,  3, 10])