Python/Numpy/Boolean Indexing: For each True value in array, modify the next 2 elements

I have a numpy array

array = np.array([5,100,100,100,5,5,100,100,100,5])

I create a mask with boolean indexing like so:

mask = (array < 30)

This gives a mask like

[ True False False False  True  True False False False  True]

I can get the indices of the True values in the mask with

indices = np.where(mask)[0]

This gives

[0 4 5 9]

For every True value in the mask, I would like to modify the next 2 elements to also be True.

I can do this with a for loop like so:

for i in indices:
    mask[i:i+3] = True

Is there a more numpythonic approach to this without using a for loop?

Desired mask output:

[ True  True  True False  True  True  True  True False  True]

The main priority here is performance.

Solution

You can use np.flatnonzero to simplify the getting of indices. Then you can add np.arange(3) to each one:

ind = np.flatnonzero(mask)[:, None] + np.arange(3)

The only caveat is that your index may contain a couple of out-of-bounds elements. You can trim them with a mask or np.clip:

ind[ind >= mask.size] = mask.size - 1

You can then apply the index directly, since numpy allows arbitrary dimensions for fancy indices:

mask[ind] = True

If you have a small smear to do, you can smear the mask directly:

mask[1:] |= mask[:-1]
mask[1:] |= mask[:-1]

You'll obviously have to put this in a loop if the smear amount is arbitrary, but you can optimize it by stepping in powers of two.

I call the operation mask[1:] |= mask[:-1] smearing because it expands the size of any group of True elements to the right by one, as if you smeared the ink with your finger. To smear an arbitrary amount n:

s = 1
while s <= n:
    mask[s:] |= mask[:-s]
    s *= 2
s = n - s // 2
if s:
    mask[s:] |= mask[:-s]