3

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

ZR-
  • 809
  • 1
  • 4
  • 12
  • 1
    Possible duplicate of https://stackoverflow.com/questions/13635032/what-is-the-inverse-function-of-zip-in-python – tripleee Aug 11 '21 at 16:16

2 Answers2

7

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]
quamrana
  • 37,849
  • 12
  • 53
  • 71
6

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

Green Cloak Guy
  • 23,793
  • 4
  • 33
  • 53
  • 1
    If `plotvex` doesn't return a 2-tuple, another solution would be to do `pstat1, ptset1, *_ = zip(*[plotvex(alpha, beta, j) for j in range(5)])` – Alex Waygood Aug 11 '21 at 16:25
  • Thanks for the answer, is there a way I can directly return `pstat1` and `ptset1` as list? – ZR- Aug 11 '21 at 16:29
  • 1
    @ZR- if you need to you can add another comprehension on the outside to convert tuple to list - `[list(tup) for tup in zip(...)]`, and it shouldn't hurt anything – Green Cloak Guy Aug 11 '21 at 16:36