J. Dionisio
J. Dionisio

Reputation: 170

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

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?

Upvotes: 1

Views: 103

Answers (2)

Riccardo Bucco
Riccardo Bucco

Reputation: 15384

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

Upvotes: 6

wasif
wasif

Reputation: 15498

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.

Upvotes: 3

Related Questions