Select NumPy Values Around Index

Question

I have two NumPy arrays:

import numpy as np

m = 3
x = np.array([1, 0, 0, np.inf, 0, 0, 1, 1, 2, np.inf, np.inf, 2, 3, 4, 5, 6, 7, 8, 9, 10])
y = np.arange(x.shape[0]-m+1)

Let's say that where ever there is an np.inf in x, that index position is called i. For each i, I want to set the values of y[i-m+1:i+m] = np.inf. So, after the replacement, y should look like:

array([0, np.inf, np.inf, np.inf, np.inf, np.inf, 6, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, 13, 14, 15, 16, 17])

This should also work when the value of m is either increased or decreased.

Answer 1

The easiest solution I can think of is to use the np.convolve function to dilate a mask. This can be done as follows:

mask = np.convolve(x==np.inf, [True]*(m*2-1), mode='same')
y[mask[:-m+1]] = np.inf

Answer 2

Here's one approach defineing an integer mask to index y and using broadcasting:

m = 3
x = np.array([1, 0, 0, np.inf, 0, 0, 1, 1, 2, np.inf, np.inf, 2, 3, 4, 5, 6, 7, 8, 9, 10])
y = np.arange(x.shape[0]-m+1).astype(float)

---

i = np.flatnonzero(x == np.inf)
y[(i + np.arange(-m+1,m)[:,None]).ravel('F')] = np.inf

print(y)
array([ 0., inf, inf, inf, inf, inf,  6., inf, inf, inf, inf, inf, inf,
       13., 14., 15., 16., 17.])