CODEWITHSUNDEEP
pythonlistlist-comprehension
James Schinner
James Schinner

Reputation: 1579

Python: If statement inside list comprehension on a generator

Python 3.6

Consider this code:

from itertools import groupby

result = [list(group) for key, group in groupby(range(5,15), key= lambda x: str(x)[0])]

print(result)

outputs:

[[5], [6], [7], [8], [9], [10, 11, 12, 13, 14]]

Can I filter out the lists with len < 2 inside the list comprehension?

Update:

Due to the two excellent answers given. I felt it might be worth a bench mark

import timeit

t1 = timeit.timeit('''
from itertools import groupby
result = [group_list for group_list in (list(group) for key, group in groupby(range(5,15), key= lambda x: str(x)[0])) if len(group_list) >= 2]
''', number=1000000)
print(t1)

t2 = timeit.timeit('''
from itertools import groupby
list(filter(lambda group: len(group) >= 2, map(lambda key_group: list(key_group[1]),groupby(range(5,15), key=lambda x: str(x)[0]))))
''', number=1000000)
print(t2)

Results:

8.74591397369441
9.647086477861325

Looks like the list comprehension has an edge.

Upvotes: 1

Views: 127

Answers (1)

rampion
rampion

Reputation: 89053

Yes

A list comprehension consists of brackets containing an expression followed by a for clause, then zero or more for or if clauses. The result will be a new list resulting from evaluating the expression in the context of the for and if clauses which follow it. For example, this listcomp combines the elements of two lists if they are not equal:

>>> [(x, y) for x in [1,2,3] for y in [3,1,4] if x != y]
[(1, 3), (1, 4), (2, 3), (2, 1), (2, 4), (3, 1), (3, 4)]

and it’s equivalent to:

>>> combs = []
>>> for x in [1,2,3]:
...     for y in [3,1,4]:
...         if x != y:
...             combs.append((x, y))
...
>>> combs
[(1, 3), (1, 4), (2, 3), (2, 1), (2, 4), (3, 1), (3, 4)]

Note how the order of the for and if statements is the same in both these snippets.

Since calling list(group) twice doesn't work in your particular example (as it consumes the generator yielded by groupby), you can introduce a temporary variable in your list comprehension by using a generator expression:

>>> [group_list for group_list in (list(group) for key, group in groupby(range(5,15), key= lambda x: str(x)[0])) if len(group_list) >= 2]
[[10, 11, 12, 13, 14]]

Alternately, using filter, map, and list:

>>> list(\
...   filter(lambda group: len(group) >= 2,\
...     map(lambda key_group: list(key_group[1]),\
...       groupby(range(5,15), key=lambda x: str(x)[0])\
...     )\
...   )\
... )
[[10, 11, 12, 13, 14]]

Upvotes: 3

Related Questions