python - for - ¿Cómo limitar el tamaño de una comprensión?

python list comprehension tutorial

Question

Tengo una list y quiero construir (a través de una comprensión) otra lista. Me gustaría que esta nueva lista tenga un tamaño limitado, a través de una condición

El siguiente código fallará:

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

con

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

porque b aún no está definido en el momento en que se construye la comprensión.

¿Hay alguna manera de limitar el tamaño de la nueva lista en tiempo de compilación?

Nota: Podría dividir la comprensión en un bucle for con la break adecuada cuando se alcanza un contador, pero me gustaría saber si hay un mecanismo que utiliza una comprensión.

Answer 1

Puede usar itertools.count para generar un contador e itertools.takewhile para detener la iteración sobre un generador cuando el contador alcanza el entero deseado ( 3 en este caso):

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

O una idea similar al construir una construcción para levantar a StopIteration para terminar el generador. Eso es lo más cerca que llegarás a tu idea original de romper la lista de comprensión , pero no lo recomendaría como la mejor práctica:

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

Ejemplos:

>>> 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]

Answer 2

Puede usar una expresión de generador para hacer el filtrado, luego use islice() para limitar el número de iteraciones:

from itertools import islice

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

Esto garantiza que no haga más trabajo de lo necesario para producir esos 3 elementos.

Tenga en cuenta que ya no tiene sentido usar una lista de comprensión aquí; una lista de comprensión no se puede romper, está bloqueado para iterar hasta el final.

Answer 3

use enumerate:

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

Answer 4

itertools.slice es la forma natural de extraer n elementos de un generador.

Pero también puede implementarlo usted mismo usando una función de ayuda. Al igual que el itertools.slice itertools.slice , StopIteration para limitar el número de elementos StopIteration .

Esto es más adaptable porque le permite especificar la lógica si n es mayor que la cantidad de elementos en su generador.

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]

Answer 5

@Martijn Pieters tiene razón en que itertools.islice es la mejor manera de resolver esto. Sin embargo, si no te importa una biblioteca adicional (externa), puedes usar iteration_utilities que envuelve muchos de estos itertools y sus aplicaciones (y algunos adicionales). Podría hacer esto un poco más fácil, al menos si te gusta la programación funcional:

>>> 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]

La clase iteration_utilities.Iterable utiliza generadores internamente, por lo que solo procesará tantos elementos como sea necesario hasta que llame a cualquiera de los as_* (o get_* ).

^{Descargo de responsabilidad: soy el autor de la biblioteca iteration_utilities .}