McLeodx McLeodx - 1 year ago 88
Python Question

Using if for multiple input lists in Python list comprehension

How do you use an if statement in a list comprehension when there are multiple input lists. Here is the code that I'm using and the error that I'm getting:

(I get that it's not able to apply modulus to a list, but not sure how to specifically reference the x in each of the lists as it iterates through them)

a = [1,2,3]
b = [4,5,6]

nums = [x**2 for x in (a,b) if x%2==0]

TypeError: unsupported operand type(s) for %: 'list' and 'int'

Answer Source

This isn't a cause with the if statement, the issue here is with x in (a, b). When that executes, x takes on a list value (first a, then b) and then Python will try try to execute your if condition on it, an expression of the form:

[1, 2, 3] % 2

is done, which obviously isn't allowed.

Instead, use chain from itertools to chain both lists together and make x take values from them:

a = [1,2,3]
b = [4,5,6]

nums = [x**2 for x in chain(a,b) if x%2==0]
[4, 16, 36]

If you're using Python >= 3.5 you could also unpack in the list literal []:

nums = [x**2 for x in [*a, *b] if x%2==0]