CODEWITHSUNDEEP
pythonlist-comprehension
ZR-
ZR-

Reputation: 819

Can I combine 2 list comprehensions in a single line?

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:)

Upvotes: 3

Views: 277

Answers (2)

Green Cloak Guy
Green Cloak Guy

Reputation: 24691

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

Upvotes: 6

quamrana
quamrana

Reputation: 39354

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]

Upvotes: 7

Related Questions