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

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!

CodePudding user response：

I would approach this using a generator that is quite efficient:

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

# [0, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8, 9, 10, 10, 11, 12, 12, 13, 14, 14, 15, 16, 16, 17, 18, 18, 19, 20]

If you really want a one-liner (not strictly a list comprehension):

from itertools import chain

n = 20
list(chain.from_iterable([i, i 1] if i%2 else [i] for i in range(n)))

# [0, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8, 9, 10, 10, 11, 12, 12, 13, 14, 14, 15, 16, 16, 17, 18, 18, 19, 20]

CodePudding user response：

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

CodePudding user response：

I do not really understand what you mean but if you wanted to create a list of even numbers with a for loop you could just do:

n = 20
nums = []
for i in range(n):
  if i % 2 == 0:
    nums.append(i)

If an integer is odd and you add one, it becomes even again. That's what makes your question a bit confusing...