Как ограничить размер понимания?

У меня есть 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, когда счетчик достигнут, но я хотел бы знать, есть ли механизм, который использует понимание.

Ответ 1

Вы можете использовать выражение генератора для фильтрации, а затем использовать islice() для ограничения числа итераций:

from itertools import islice

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

Это гарантирует, что вы не будете делать больше работы, чем вам нужно, чтобы создать эти 3 элемента.

Обратите внимание, что больше нет смысла использовать понимание списка; понимание списка не может быть нарушено, вы заперты в итерации до конца.

Ответ 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.}

Ответ 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]

Ответ 4

использовать перечисление:

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