CODEWITHSUNDEEP
python
hughdbrown
hughdbrown

Reputation: 49013

Generate two lists at once

Background

The algorithm manipulates financial analytics. There are multiple lists of the same size and they are filtered into other lists for analysis. I am doing the same filtering on different by parallel lists. I could set it up so that a1,b1,c2 occur as a tuple in a list but then the analytics have to stripe the tuples the other way to do analysis (regression of one list against the other, beta, etc.).

What I want to do

I want to generate two different lists based on a third list:

>>> a = list(range(10))
>>> b = list(range(10,20))
>>> c = list(i & 1 for i in range(10))
>>>
>>> aprime = [a1 for a1, c1 in zip(a,c) if c1 == 0]
>>> bprime = [b1 for b1, c1 in zip(b,c) if c1 == 0]
>>> aprime
[0, 2, 4, 6, 8]
>>> bprime
[10, 12, 14, 16, 18]

It seems there should be a pythonic/functional programming/itertools way to create the two lists and iterate over the three lists only once. Something like:

aprime, bprime = [a1, b1 for a1, b1, c1 in zip(a,b,c) if c1 == 0]

But of course this generates a syntax error.

The question

Is there a pythonic way?

Micro-optimization shootout

The ugly but pythonic-to-the-max one-liner edges out the "just use a for-loop" solution and my original code in the ever popular timeit cage match:

>>> import timeit
>>> timeit.timeit("z2(a,b,c)", "n=100;a = list(range(n)); b = list(range(10,10+n)); c = list(i & 1 for i in range(n));\ndef z2(a,b,c):\n\treturn zip(*[(a1,b1) for a1,b1,c1 in zip(a,b,c) if c1==0])\n")
26.977873025761482
>>> timeit.timeit("z2(a,b,c)", "n=100;a = list(range(n)); b = list(range(10,10+n)); c = list(i & 1 for i in range(n));\ndef z2(a,b,c):\n\taprime, bprime = [], [];\n\tfor a1, b1, c1 in zip(a, b, c):\n\t\tif c1 == 0:\n\t\t\taprime.append(a1);  bprime.append(b1);\n\treturn aprime, bprime\n")
32.232914169258947
>>> timeit.timeit("z2(a,b,c)", "n=100;a = list(range(n)); b = list(range(10,10+n)); c = list(i & 1 for i in range(n));\ndef z2(a,b,c):\n\treturn [a1 for a1, c1 in zip(a,c) if c1 == 0], [b1 for b1, c1 in zip(b,c) if c1 == 0]\n")
32.37302275847901

Upvotes: 2

Views: 2982

Answers (3)

sepp2k
sepp2k

Reputation: 370192

This might win the ugliest code award, but it works in one line:

aprime, bprime = zip(*[(a1,b1) for a1,b1,c1 in zip(a,b,c) if c1==0])

Upvotes: 5

Isaiah
Isaiah

Reputation: 4309

Just use a for loop:

aprime = []
bprime = []
for a1, b1, c1 in zip(a, b, c):
    if c1 == 0:
        aprime.append(a1) 
        bprime.append(b1)

Upvotes: 4

Sam DeFabbia-Kane
Sam DeFabbia-Kane

Reputation: 2609

There's no way to create multiple lists at a time with list comprehensions--if you only want to iterate once you're going to need to do it some other way--possible with a loop.

You could use a list comprehension to create a list of tuples, with the first element belonging to one list, the second to the other. But if you do want them as separate lists, you're going to have to use another operation to split it, anyway.

Upvotes: 0

Related Questions