I have an array myA
like this:
array([ 7, 4, 5, 8, 3, 10])
If I want to replace all values that are larger than a value val
by 0, I can simply do:
myA[myA > val] = 0
which gives me the desired output (for val = 5
):
array([0, 4, 5, 0, 3, 0])
However, my goal is to replace not all but only the first n
elements of this array that are larger than a value val
.
So, if n = 2
my desired outcome would look like this (10
is the third element and should therefore not been replaced):
array([ 0, 4, 5, 0, 3, 10])
A straightforward implementation would be:
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
That works but maybe someone can can up with a smart way of masking!?
The following should work:
myA[(myA > val).nonzero()[0][:2]] = 0
since nonzero will return the indexes where the boolean array myA > val
is non zero e.g. True
.
For example:
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])
Collected from the Internet
Please contact [email protected] to delete if infringement.
Comments