python list comprehension to produce two values in one iteration

PythonListList Comprehension

Python Problem Overview


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?

Python Solutions


Solution 1 - Python

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

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

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]

Solution 4 - Python

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]

Solution 5 - Python

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

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

Solution 6 - Python

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.

Solution 7 - Python

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]

Solution 8 - Python

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

Solution 9 - Python

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]

Solution 10 - Python

Another option:

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

Solution 11 - Python

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

Solution 12 - Python

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

Solution 13 - Python

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

Attributions

All content for this solution is sourced from the original question on Stackoverflow.

The content on this page is licensed under the Attribution-ShareAlike 4.0 International (CC BY-SA 4.0) license.

Content TypeOriginal AuthorOriginal Content on Stackoverflow
Questionuser1629366View Question on Stackoverflow
Solution 1 - PythonAshwini ChaudharyView Answer on Stackoverflow
Solution 2 - PythonYannView Answer on Stackoverflow
Solution 3 - PythonJonathan MayerView Answer on Stackoverflow
Solution 4 - PythonMartijn PietersView Answer on Stackoverflow
Solution 5 - PythonJoran BeasleyView Answer on Stackoverflow
Solution 6 - PythonJohn HallView Answer on Stackoverflow
Solution 7 - PythoniruvarView Answer on Stackoverflow
Solution 8 - PythonHielke WalingaView Answer on Stackoverflow
Solution 9 - PythonDMizView Answer on Stackoverflow
Solution 10 - PythonMatthew PlourdeView Answer on Stackoverflow
Solution 11 - PythonSupriya KView Answer on Stackoverflow
Solution 12 - PythonntgView Answer on Stackoverflow
Solution 13 - Pythonrassa45View Answer on Stackoverflow