Whats the most Pythonic way to make the next n items in an array True, wherever there is a True value in the array

Viewed 107

I have an array of booleans

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

and wherever there is a True value, I want to make the next n elements of the array True. If in this case I chose n=2, this array would become

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

Whats the most pythonic way to do this, without using for loops?

3 Answers

A quick way would be using np.convolve:

rng = np.random.default_rng()
a = rng.integers(0,4,20)<1
a
# array([False, False, False, False, False, False, False, False,  True,
#        False, False,  True, False, False, False, False,  True, False,
#        False, False])
n = 3
np.convolve(a,np.ones(n+1,bool),"full")[:-n]
# array([False, False, False, False, False, False, False, False,  True,
#        True,  True,  True,  True,  True,  True, False,  True,  True,
#        True,  True])

I'll assume you mean list and not array.array or numpy.ndarray.

So for your example:

foo = [False, False, True, True, False, True, False, False, False]
n = 2

The most pythonic way is - in my opinion - this one-liner list comprehension

[any(foo[max(0,i-n):i+1]) for i in range(len(foo))]

# [False, False, True, True, True, True, True, True, False]

With numpy, you can use stride tricks as_strided:

import numpy as np
from numpy.lib.stride_tricks import as_strided
 
s = np.array(
    [False, False, True, True, False, True, False, False, False])
n = 2

s2 = s[:]
s2[n:] = as_strided(s, (len(s) - (n), n+1), s.strides * 2).max(axis=1)

print (s2)
# array([False, False,  True,  True,  True,  True,  True,  True, False])

This is pretty fast.


If you have pandas, this is a rolling operation:

import pandas as pd

s = pd.Series([False, False, True, True, False, True, False, False, False])
n = 2

s.rolling(window=n+1).max().fillna(s).astype(bool)

0    False
1    False
2     True
3     True
4     True
5     True
6     True
7     True
8    False
dtype: bool
Related