如何限制推导式的大小？

Question

我有一个

list

并想构建（通过理解）另一个列表。我希望通过一个条件来限制这个新列表的大小

以下代码将失败：

a = [1, 2, 1, 2, 1, 2]
b = [i for i in a if i == 1 and len(b) < 3]

与

Traceback (most recent call last):
  File "compr.py", line 2, in <module>
    b = [i for i in a if i == 1 and len(b) < 3]
  File "compr.py", line 2, in <listcomp>
    b = [i for i in a if i == 1 and len(b) < 3]
NameError: name 'b' is not defined

因为在构建理解时

b

尚未定义。

有没有办法在构建时限制新列表的大小？

注意：当达到计数器时，我可以使用适当的

for

将理解分解为

break

循环，但我想知道是否存在使用理解的机制。

Answer 1

islice()

来限制迭代次数：

from itertools import islice

filtered = (i for i in a if i == 1)
b = list(islice(filtered, 3))

这可确保您不会做超出生成这 3 个元素所需的工作。

请注意，这里不再使用列表理解了；列表理解无法被打破，你被锁定迭代到最后。

Answer 2

@Martijn Pieters

完全正确，itertools.islice 是解决此问题的最佳方法。但是，如果您不介意额外的（外部）库，您可以使用

iteration_utilities

，它包含很多这些 

itertools 及其应用程序（以及一些其他应用程序）。它可以让这变得更容易一些，至少如果你喜欢函数式编程的话：

>>> from iteration_utilities import Iterable

>>> Iterable([1, 2, 1, 2, 1, 2]).filter((1).__eq__)[:2].as_list()
[1, 1]

>>> (Iterable([1, 2, 1, 2, 1, 2])
...          .filter((1).__eq__)   # like "if item == 1"
...          [:2]                  # like "islice(iterable, 2)"
...          .as_list())           # like "list(iterable)"
[1, 1]

iteration_utilities.Iterable

类在内部使用生成器，因此它只会处理所需数量的项目，直到您调用任何

as_*（或

get_*

）方法。

---

免责声明：我是

^{iteration_utilities库}

的作者。

Answer 3

itertools.count

生成计数器，并在计数器达到所需整数时使用

itertools.takewhile

停止对生成器的迭代（在本例中为

3

）：

from itertools import count, takewhile
c = count()
b = list(takewhile(lambda x: next(c) < 3, (i for i in a if i == 1)))

或者类似的想法构建一个构造来引发 

StopIteration

来终止生成器。这是您最接近“打破列表理解”的原始想法，但我不会推荐它作为最佳实践：

c = count()
b = list(i if next(c) < 3 else next(iter([])) for i in a if i == 1)

示例：

>>> a = [1,2,1,4,1,1,1,1]

>>> c = count()
>>> list(takewhile(lambda x: next(c) < 3, (i for i in a if i == 1)))
[1, 1, 1]

>>> c = count()
>>> list(i if next(c) < 3 else next(iter([])) for i in a if i == 1)
[1, 1, 1]

相同的解决方案，只是没有

islice

Answer 4

filtered = (i for i in a if i == 1)
b = [filtered.next() for j in range(3)]

顺便说一句，请注意生成器是否为空或者生成器少于 3 个 - 你会得到

StopIteration Exception

。

为了防止这种情况，您可能需要使用带有默认值的 next() 。例如： b = [next(filtered, None) for j in range(3)]

如果您不希望列表中出现“无”：

b = [i for i in b if i is not None]

Answer 5

是从生成器中提取 n

 项的自然方法。

但是您也可以使用辅助函数自己实现这一点。就像 itertools.slice

伪代码

一样，我们捕获

StopIteration

来限制生成的项目数量。 这更具适应性，因为它允许您在

n

大于生成器中的项目数时指定逻辑。

def take_n(gen, n):
    for _ in range(n):
        try:
            yield next(gen)
        except StopIteration:
            break

g = (i**2 for i in range(5))
res = list(take_n(g, 20))

print(res)

[0, 1, 4, 9, 16]

a = [1, 2, 1, 2, 1, 2]

Answer 6

我认为这会创建一个完整的列表理解（评估原始列表中的每个元素），然后对其进行切片。它在长列表中可能不会有很好的性能，但很容易阅读，并且写入速度非常快。

使用枚举：

Answer 7