insert element in a list after every 2 element

Question

I have this code:

l = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]

def add(l, n):
    lret = l.copy()
    i = n
    while i <= len(lret):
      lret.insert(i, 0.0)
      i += (n+1)
    return lret

lret = add(l,2)

print(lret)

Is a way to make the add() function a one-line function?

Answer 1

You can use zip() to build tuples that contain n elements from the original list with an added element, and then use chain.from_iterable() to flatten those tuples into the final list:

from itertools import repeat, chain
def add(lst, n):
    return list(chain.from_iterable(zip(*(lst[s::n] for s in range(n)), repeat(0.0))))

This outputs:

[1.0, 2.0, 0.0, 3.0, 4.0, 0.0, 5.0, 6.0, 0.0]