Modifying a chunk of arrays in Numpy Python

Viewed 66

I want to write a function where if an array is equal to or less than 0 it will replace the values after (and including) that index with 0. The numbers have the 7th element less than 0 being -1, thus elements including the 7th element to 10th element values will be replaced by 0. How could I configure the the np.where function so it gives the expected result values?

Sample input:

numbers = np.array([87., 4., 4.1,  8.5, 10 , 20, 22.3, -1., 1., 60., -4.])

Expected output:

[87., 4., 4.1,  8.5, 10 , 20, 22.3, 0., 0., 0., 0.]
3 Answers

You can use argmax to find out the first indice that's smaller than 0, and then modify all values after that indice to 0:

a = np.array([87. ,4., 4.1,  8.5, 10 , 20, 22.3, -1, 1, 60, -4])
if (a <= 0).any(): a[(a <= 0).argmax():] = 0

a
# [87.   4.   4.1  8.5 10.  20.  22.3  0.   0.   0.   0. ]

Or more efficiently, you can check the value under the candidate index only:

def zero_after_native(a):
    candidate = (a <= 0).argmax()
    if a[candidate] <= 0:
        a[candidate:] = 0

a = np.array([87. ,4., 4.1,  8.5, 10 , 20, 22.3, -1, 1, 60, -4])
zero_after_native(a)

a
# [87.   4.   4.1  8.5 10.  20.  22.3  0.   0.   0.   0. ]

Another solution is:

Numbers[np.where(Numbers<=0)[0][0]:]=0

or equally:

Numbers[np.flatnonzero(Numbers<=0)[0]:]=0

or equally:

Numbers[np.nonzero(Numbers<=0)[0][0]:]=0

EDIT:

The above solutions work assuming you have at least one non-positive number in your array. If it is not guaranteed, you can use an if statement before them like:

if (Numbers<=0).any(): 
  Numbers[np.where(Numbers<=0)[0][0]:]=0

output:

array([87. ,  4. ,  4.1,  8.5, 10. , 20. , 22.3,  0. ,  0. ,  0. ,  0. ])

First off, don't use np.where. For what you have here, it's just a wrapper for np.nonzero, which you don't need because this can all be done using masks.

To correctly handle the case without any negatives, I'd recommend finding the last positive index rather than the first negative index, and setting everything to zero after that:

def zero_tail(a):
    a[a.size - (a >= 0)[::-1].argmax():] = 0

To check:

>>> a = np.arange(9, -5, -2)
>>> zero_tail(a)
>>> a

>>> a = np.arange(1, 10)
>>> zero_tail(a)
>>> a

If you absolutely wanted to use np.nonzero instead of np.argmax, you could do something similar (I'd actually recommend np.flatnonzero to avoid an extra layer of unpacking):

def zero_tail(a):
    a[a.size - np.flatnonzero((a >= 0)[::-1])[0]:] = 0
Related