CODEWITHSUNDEEP
pythonlistiteratorlist-comprehension
Jeremy
Jeremy

Reputation: 129

Generate multiple values in each iteration of a list comprehension

I'd like to generate a list, say, [0, 0^2, 1, 1^2, ... 9, 9^2] in python. However, [(i,i**2) for i in range(10)] returns a list of tuples and [*(i,i**2) for i in range(10)] doesn't work.

Is there any pythonic way to create such a list?

Upvotes: 1

Views: 516

Answers (1)

yatu
yatu

Reputation: 88246

You could add another level of iteration in the list comprehension to get a flat list:

[i**exp for i in range(10) for exp in [1,2]]
# [0, 0, 1, 1, 2, 4, 3, 9, 4, 16, 5, 25, 6, 36, 7, 49, 8, 64, 9, 81]

Or you could use itertools.chain:

from itertools import chain

list(chain.from_iterable((i,i**2) for i in range(10)))
# [0, 0, 1, 1, 2, 4, 3, 9, 4, 16, 5, 25, 6, 36, 7, 49, 8, 64, 9, 81]

Though I'd consider numpy for such task:

(np.arange(10)**np.arange(1,3)[:,None]).ravel('F')

array([ 0,  0,  1,  1,  2,  4,  3,  9,  4, 16,  5, 25,  6, 36,  7, 49,  8,
   64,  9, 81], dtype=int32)

Upvotes: 3

Related Questions