如何限制理解的大小？答案

【问题标题】：How to limit the size of a comprehension?如何限制理解的大小？
【发布时间】：2021-11-06 00:11:06
【问题描述】：

我有一个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 在构建理解时尚未定义。

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

注意：当达到计数器时，我可以使用正确的 break 将理解分解为 for 循环，但我想知道是否存在使用理解的机制。

【问题讨论】：

标签： python list-comprehension

【解决方案1】：

您可以使用生成器表达式进行过滤，然后使用islice() 限制迭代次数：

from itertools import islice

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

这可确保您不会做比生成这 3 个元素更多的工作。

请注意，这里不再使用列表推导式了；列表推导不能被打破，你被锁定在迭代到最后。

【讨论】：

[1/i for i in range(-5, 5)] 确实爆发并且没有迭代到最后。
@StefanPochmann：它引发了一个异常，这与break 语句不同。最后，您根本没有列表结果。
我不清楚你的意思是break 声明，这个词可以用更一般的方式理解。 For example不久前你说“[return]跳出循环”。在任何情况下，迭代都不会结束。此外，没有列表结果甚至都不是问题。考虑reciprocals = [1/x for x in a]，我认为这是合理的代码，如果a 包含零，那么可能需要ZeroDivisionError 而不需要列表。
这是一个关于如何限制由列表理解生成的列表大小的问题。这意味着您仍需要列表结果。

【解决方案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_*）方法。

^{免责声明：我是the iteration_utilities library的作者。}

【讨论】：

这是一个非常好的库，谢谢（仍在阅读文档以掌握众多功能）
我可以建议更改项目主页的第一个链接：iteration-utilities.readthedocs.io/en/latest？

【解决方案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 以终止生成器。这是您最接近 breaking the list comprehension 最初想法的方式，但我不建议将其作为最佳实践：

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]

【讨论】：

与其他答案相比，这有什么优势？
@jpmc26 对于这个确切的目的，我认为它并不比 Martijn 的解决方案更好，但它更通用，因为终止生成器的条件可以是任何东西，而不仅仅是计数器。 OP还专门询问了列表理解，这是最接近的有效语法
很公平。谢谢。由于您在其他答案之后很好地发布了此内容，因此您可能希望在您的答案中考虑灵活性优势。

【解决方案4】：

使用枚举：

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

【讨论】：

那是完全错误的。首先，这将丢弃除a 的前 3 项之外的所有内容（问题想限制b 而不是a），它将处理整个迭代。找到第三项后它不会停止。此后它只会丢弃所有内容（但它会坚持检查n==1 and i < 3）。

【解决方案5】：

itertools.slice 是从生成器中提取 n 个项目的自然方式。

但您也可以使用辅助函数自行实现。就像 itertools.slice pseudo-code 一样，我们捕获 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]

【讨论】：

【解决方案6】：

同样的解决方案只是没有islice:

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]

【讨论】：