Can I combine 2 list comprehensions in a single line?

Viewed 91

Here's part of my Python code:

pstat1 = [plotvex(alpha,beta,j)[0] for j in range(5)]
ptset1 = [plotvex(alpha,beta,j)[1] for j in range(5)]

where plotvex is a function that returns 2 items. I want to generate two lists pstat1 and ptset1 using list comprehension, but I wonder is there a way I don't need to call the function twice? Thanks:)

2 Answers

You are quite right that you don't want to call the plotvex() function twice for each set of parameters.

So, just call it once and then generate pstat1 and pstat2 later:

pv = [plotvex(alpha,beta,j) for j in range(5)]
pstat1 = [item[0] for item in pv]
ptset1 = [item[1] for item in pv]

Assuming plotvex() returns a 2-tuple exactly*, this should work:

pstat1, ptset1 = zip(*[plotvex(alpha, beta, j) for j in range(5)])

zip(*iterable_of_iterables) is a common idiom to 'rotate' a list of lists from being vertical to being horizontal. So instead of a list of 2-tuples, [plotvex(alpha, beta, j) for j in range(5)] will become two lists of singles, one list from each half of the tuples.

* here is the argument-unpacking operator.


*if it returns more than a 2-tuple, then just do plotvex(alpha, beta, j)[:2] instead to take the first two elements

Related