How can I create more than one array at the same time using a generator?

Question

Let's assume we have 3 arrays, index, a and b. How can I create arrays c and d just passing through index once?

c = [a[i] for i in index]
d = [b[i] for i in index]

Is there a way to create these arrays with a sigle generator?

Answer 1

You can use the zip function:

c, d = zip(*((a[i], b[i]) for i in index))

If you want c and d to be lists you can use map :

c, d = map(list, zip(*((a[i], b[i]) for i in index)))

If you want something longer (but maybe clearer), you could build a generator:

def g(a, b, index):
    for i in index:
        yield a[i], b[i]

c, d = zip(*g(a, b, index))

Answer 2

I will use zip with tuple expansion using *

c,d = zip(*((a[i],b[i]) for i in index))

Here it expands a[i] and b[i] in pairs using zip from generator expression.