Conditionally adding 1 or 2 items into a list using a list comprehension

Question

Is there a way to insert one or two items into a list using a list comprehension depending on some condition using only one for loop?

For example if I wanted to iterate over n numbers, if a number is even then insert it and if a number is odd then insert it and also insert it plus 1.

So far all I've got is by using a nested list comprehension with three for loops total when this should be done with only one for loop:

new_lst = [num for sublist in [[i] if i%2 == 0 else [i,i+1] for i in range(n)] for num in sublist]

The normal way with one for loop would be like so:

new_list = []
for i in range(n):
    new_lst.append(i)
    if i % 2 == 1:
        new_lst.append(i+1)

Any help would be much appreciated, thanks!

Answer 1

I would approach this using two separate list comprehensions. This should improve time complexity slightly, down to O(N) + O(N/2) which simplifies to O(N).

from itertools import chain
from timeit import timeit


def gen(n):
    for i in range(n):
        yield i
        if i % 2:
            yield i + 1


def gen2(n):
    return [i for i in range(n + 1)] + [i for i in range(2, n, 2)]


def gen2_in_order(n):
    return sorted([i for i in range(n + 1)] + [i for i in range(2, n, 2)])


def gen3(n):
    return list(chain.from_iterable([i, i + 1] if i % 2 else [i] for i in range(n)))


def gen4(n):
    return [num for i in range(n) for num in range(i, i + 1 + i % 2)]


assert list(gen(1000)) == gen2_in_order(1000) == gen3(1000) == gen4(1000)

print('generator:           ', timeit('gen(n)', globals=globals(), setup='n=100'))
print('generator (list):    ', timeit('list(gen(n))', globals=globals(), setup='n=100'))
print('list comp:           ', timeit('gen2(n)', globals=globals(), setup='n=100'))
print('list comp (sorted):  ', timeit('gen2_in_order(n)', globals=globals(), setup='n=100'))
print('from_iter:           ', timeit('gen3(n)', globals=globals(), setup='n=100'))
print('list comp (single):  ', timeit('gen4(n)', globals=globals(), setup='n=100'))

Results on my Mac:

generator:            0.08834924991242588
generator (list):     5.09551537502557
list comp:            2.2710815421305597
list comp (sorted):   3.3015330410562456
from_iter:            7.440466790925711
list comp (single):   12.89962362498045