python list comprehension to produce two values in one iteration
PythonListList ComprehensionPython 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)]