Numpy fastest way to create mask array around indexes-CodePudding

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

CodePudding user response：

Find the elements larger than 2 then set the elements around them to True:

a = np.array([0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 2, 0])

N = 3
mask = a > 2
center = np.where(mask)[0]
mask[np.maximum(np.ravel(center - np.arange(1, 1   N).reshape(-1, 1)), 0)] = True
mask[np.minimum(np.ravel(center   np.arange(1, 1   N).reshape(-1, 1)), len(a)-1)] = True

Thanks @Michael Szczesny for pointing out the edge case. The maximum and minimum ensures that the index would not (unintentionally) go out of bounds.

CodePudding user response：

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