使用列表理解有条件地将 1 或 2 个项目添加到列表中

Question

有没有一种方法可以使用列表推导式根据某些条件仅使用一个 for 循环将一个或两个项目插入到列表中？

例如，如果我想遍历 n 个数字，如果一个数字是偶数则插入它，如果一个数字是奇数则插入它并插入它加 1。

到目前为止，我所得到的只是使用嵌套列表理解和三个 for 循环，而这应该只用一个 for 循环来完成：

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

一个for循环的正常方式是这样的：

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

非常感谢任何帮助，谢谢！

Answer 1

我会使用非常高效的生成器来解决这个问题：

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]

如果你真的想要一个one-liner（不是严格意义上的列表理解）：

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]

Answer 2

我会使用两个单独的 list 理解来解决这个问题。这应该会略微提高时间复杂度，降低到 O(N) + O(N/2)，从而简化为 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'))

我的 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

使用列表理解有条件地将 1 或 2 个项目添加到列表中

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

python

list

list-comprehension