If True: turn following values to True in list of booleans until nth position after True is reached

Question

I have a list of booleans

l = [False, False, False, True, False, False, False]

that I want to turn into

l_new = [False, False, False, True, True, True, False]

That means, whenever there is a True in my list I want to switch the two (for example) following values to true. My solution is

def lagged_effect(l, lag):
    l_new = []
    L_iter = iter(l)
    for elem in L_iter:
        if elem == True:
            l_new.extend([True]*lag)
            if lag == 1:
                next(L_iter)
            if lag == 2:
                next(L_iter)
                next(L_iter)
            if lag == 3:
                next(L_iter)
                next(L_iter)
                next(L_iter)
            if lag == 4:
                next(L_iter)
                next(L_iter)
                next(L_iter)
                next(L_iter)
            if lag > 4:
                print("not defined")
        if elem == False:
            l_new.append(False)
    return l_new
    print(l_new)
lagged_effect(l, lag=2)

Since I want to implement this more often, I was wondering if there might be a more compact and efficient solution. Especially the next implementation annoys me.

Answer 1

You can write a custom function:

l = [False, False, False, True, False, False, False]

def lag_true(l, lag=2):
    out = []
    last = 0
    for v in l:
        if v:                # if True, set up last to lag value
            last = lag
            out.append(True)
        elif last:           # if last (>0, implicit)
            out.append(True) # add True
            last -= 1        # decrease last
        else:                # we have a False and are far from last True
            out.append(False)
    return out

lag_true(l)
# [False, False, False, True, True, True, False]

lag_true(l, lag=3)
# [False, False, False, True, True, True, True]

Vectorial alternative using pandas

import pandas as pd

LAG = 2

out = pd.Series(l).rolling(LAG+1, min_periods=1).max().astype(bool).to_list()
# [False, False, False, True, True, True, False]