python list comprehension to produce two values in one iteration

I want to generate a list in python as follows -

[1, 1, 2, 4, 3, 9, 4, 16, 5, 25 .....]

You would have figured out, it is nothing but n, n*n

I tried writing such a list comprehension in python as follows -

lst_gen = [i, i*i for i in range(1, 10)]

But doing this, gives a syntax error.

What would be a good way to generate the above list via list comprehension?

Solution 1:

Use itertools.chain.from_iterable:

>>> from itertools import chain
>>> list(chain.from_iterable((i, i**2) for i in xrange(1, 6)))
[1, 1, 2, 4, 3, 9, 4, 16, 5, 25]

Or you can also use a generator function:

>>> def solve(n):
...     for i in xrange(1,n+1):
...         yield i
...         yield i**2

>>> list(solve(5))
[1, 1, 2, 4, 3, 9, 4, 16, 5, 25]

Solution 2:

The question is old, but just for the curious reader, i propose another possibility: As stated on first post, you can easily make a couple (i, i**2) from a list of numbers. Then you want to flatten this couple. So just add the flatten operation in your comprehension.

[x for i in range(1, 10) for x in (i,i**2)]

Solution 3:

A little-known trick: list comprehensions can have multiple for clauses.

For example:

>>> [10*x+y for x in range(4) for y in range(3)]
[0, 1, 2, 10, 11, 12, 20, 21, 22, 30, 31, 32]

In your particular case, you could do:

>>> [x*x if y else x for x in range(5) for y in range(2)]
[0, 0, 1, 1, 2, 4, 3, 9, 4, 16]