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?

49164 次浏览

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

oh I should mention the sum probably breaks the one iteration rule :(

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]

Another option:

reduce(lambda x,y: x + [y, y*y], range(1,10), [])

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]

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.

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

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

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

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

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]

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

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]