具有两个输出的列表理解

Question

I am trying to optimize this code:

num = 10
for j in xrange(0,num):
  u[j],v[j] = rk4(du,dv,t,dt,u[j],v[j])

where u and v are input arrays and rk4() returns two values for two input values. Using list comprehension I would do something like this:

u,v=[rk4(du,dv,t,dt,u[j],v[j])) for j in range(0,num)]

The list comprehension works. But the output is in a different format. Is it possible to optimize this kind of operation using list comprehension?

Edit: The desired output would be two arrays/lists of the form

u,v = [u1,u2,u3,....],[v1,v2,v3,...]

What I get is the of the following form:

[(u1,v1),(u2,v2),(u3,v3),...]

Answer 1

It appears that you want to transform a sequence of pairs into two sequences. There is a standard idiom in Python to do this using the zip function and argument unpacking:

>>> seq_of_pairs = [('a', 1), ('b', 2), ('c', 3), ('d', 4)]
>>> u, v = zip(*seq_of_pairs)
>>> u
('a', 'b', 'c', 'd')
>>> v
(1, 2, 3, 4)

So you can use a list comprehension (or generator expression) to produce the sequence of pairs using zip , and then use that trick to extract the two sequences:

result = [ rk4(..., ui, vi) for ui, vi in zip(u, v) ]
u, v = zip(*result)

You can do u, v = map(list, zip(*result)) if you need them to be lists instead of tuples.