'Numpy fastest way to create mask array around indexes

I have a 1D array:

arr = np.array([0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 2, 0, 0, 0, 0, 0, 2, 3, 0, 0, 1, ...], dtype='uint16')

I want to create a mask array that is True anywhere that is +/- N indexes of of a value greater than 2, yielding the following (for N=3)

mask = [F, T, T, T, T, T, T, T, F, F, T, T, T, T, T, T, T, F, F, F, T, T, T, T, T, T, T, ...]

(note, I used T/F for readability's sake)

I need this to be rather fast as my actual array is millions of points long, N will likely be something like 500.

Edit: similar problem

pythonarraysalgorithmnumpy


Solution 1:[1]

You can use a convolution with numpy.convolve and an array of ones of length 2*N+1 (N on each side plus self) as mask:

np.convolve(arr>2, np.ones(2*N+1), mode='same').astyoe(bool)

Output:

array([False,  True,  True,  True,  True,  True,  True,  True, False,
       False,  True,  True,  True,  True,  True,  True,  True, False,
       False, False,  True,  True,  True,  True,  True,  True,  True])

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1 mozway

Related Questions