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

Question

我有这样一个数组myA:

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

如果我想将大于值的所有值替换val为0,我可以简单地执行:

myA[myA > val] = 0

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

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

但是,我的目标是替换不仅仅是n这个数组中大于值的第一个元素val.

所以,如果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

因为非零将返回布尔数组非零的索引,myA > val例如True.

例如:

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])