python list comprehension to produce two values in one iteration

Question:

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?

Asked By: user1629366

||

Answers:

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]
Answered By: Ashwini Chaudhary
lst_gen = sum([(i, i*i) for i in range(1, 10)],())

oh I should mention the sum probably breaks the one iteration rule 🙁

Answered By: Joran Beasley

List comprehensions generate one element at a time. Your options are, instead, to change your loop to only generate one value at a time:

[(i//2)**2 if i % 2 else i//2 for i in range(2, 20)]

or to produce tuples then flatten the list using itertools.chain.from_iterable():

from itertools import chain

list(chain.from_iterable((i, i*i) for i in range(1, 10)))

Output:

>>> [(i//2)**2 if i % 2 else i//2 for i in range(2, 20)]
[1, 1, 2, 4, 3, 9, 4, 16, 5, 25, 6, 36, 7, 49, 8, 64, 9, 81]
>>> list(chain.from_iterable((i, i*i) for i in range(1, 10)))
[1, 1, 2, 4, 3, 9, 4, 16, 5, 25, 6, 36, 7, 49, 8, 64, 9, 81]
Answered By: Martijn Pieters

Another option:

reduce(lambda x,y: x + [y, y*y], range(1,10), [])
Answered By: Matthew Plourde

Another option, might seem perverse to some

>>> from itertools import izip, tee
>>> g = xrange(1, 11)
>>> x, y = tee(g)
>>> y = (i**2 for i in y)
>>> z = izip(x, y)
>>> output = []
>>> for k in z:
...     output.extend(k)
... 
>>> print output
[1, 1, 2, 4, 3, 9, 4, 16, 5, 25, 6, 36, 7, 49, 8, 64, 9, 81, 10, 100]
Answered By: iruvar

You can create a list of lists then use reduce to join them.

print [[n,n*n] for n in range (10)]

[[0, 0], [1, 1], [2, 4], [3, 9], [4, 16], [5, 25], [6, 36], [7, 49], [8, 64], [9, 81]]

print reduce(lambda x1,x2:x1+x2,[[n,n*n] for n in range (10)])

[0, 0, 1, 1, 2, 4, 3, 9, 4, 16, 5, 25, 6, 36, 7, 49, 8, 64, 9, 81]

 print reduce(lambda x1,x2:x1+x2,[[n**e for e in range(1,4)]
 for n in range (1,10)])

[1, 1, 1, 2, 4, 8, 3, 9, 27, 4, 16, 64, 5, 25, 125, 6, 36, 216, 7, 49, 343, 8, 64, 512, 9, 81, 729]

Reduce takes a callable expression that takes two arguments and processes a sequence by starting with the first two items. The result of the last expression is then used as the first item in subsequent calls. In this case each list is added one after another to the first list in the list of lists and then that list is returned as a result.

List comprehensions implicitly call map with a lambda expression using the variable and sequence defined by the “for var in sequence” expression. The following is the same sort of thing.

map(lambda n:[n,n*n],range(1,10))

[[1, 1], [2, 4], [3, 9], [4, 16], [5, 25], [6, 36], [7, 49], [8, 64], [9, 81]]

I am unaware of a more natural python expression for reduce.

Answered By: John Hall
>>> lst_gen = [[i, i*i] for i in range(1, 10)]
>>> 
>>> lst_gen
[[1, 1], [2, 4], [3, 9], [4, 16], [5, 25], [6, 36], [7, 49], [8, 64], [9, 81]]
>>> 
>>> [num for elem in lst_gen for num in elem]
[1, 1, 2, 4, 3, 9, 4, 16, 5, 25, 6, 36, 7, 49, 8, 64, 9, 81]

Here is my reference
http://docs.python.org/2/tutorial/datastructures.html

Answered By: Supriya K

Try this two liner

lst = [[i, i*i] for i in range(10)]
[lst.extend(i) for i in lst]

Change math as necessary.

EVEN BETTER

#Change my_range to be the number you want range() function of
start = 1
my_range = 10
lst = [i/2 if i % 2 == 0 else ((i-1)/2)**2 for i in range(start *2, my_range*2 - 1)]
Answered By: rassa45

As mentioned, itertools is the way to go. Here’s how I would do it, I find it more clear:

[i if turn else i*i for i,turn in itertools.product(range(1,10), [True, False])]
Answered By: ntg

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)]
Answered By: Yann

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]
Answered By: Jonathan Mayer

Lots of tricks in this thread. Here is another using a one liner generator without imports

x = (lamdba : [[(yield i), (yield i**2)] for i in range(10)])()

EDIT:
This will raise DeprecatedWarning in Python 3.7 and SyntaxError in Python 3.8:
https://docs.python.org/dev/whatsnew/3.7.html#deprecated-python-behavior

Answered By: Hielke Walinga

A Simple one might be:

[i ** j for i in range(1, 6) for j in [1, 2]]

Which is quite explicit regarding the task and the fact that the list comp replaces two nested loops (while as other answers show, can be flattened, so to speak).
Output:

[1, 1, 2, 4, 3, 9, 4, 16, 5, 25]
Answered By: DMiz
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.