CODEWITHSUNDEEP
pythonfor-looplist-comprehension
Fosa
Fosa

Reputation: 572

Python: How to extend or append multiple elements in list comprehension format?

I'd like to get a nice neat list comprehension for this code or something similar!

extra_indices = []
for i in range(len(indices)):
    index = indices[i]
    extra_indices.extend([index, index + 1, index +2])

Thanks!

Edit* The indices are a list of integers. A list of indexes of another array.

For example if indices is [1, 52, 150] then the goal (here, this is the second time I've wanted two separate actions on continuously indexed outputs in a list comprehension)

Then extra_indices would be [1, 2, 3, 52, 53, 54, 150, 151, 152]

Upvotes: 7

Views: 2787

Answers (2)

Skycc
Skycc

Reputation: 3555

code below should do the equivalant of your code, assuming indices is a list of integer

from itertools import chain
extra_indices = list(chain(*([x,x+1,x+2] for x in indices)))

>>> indices = range(3)
>>> list(chain(*([x,x+1,x+2] for x in indices)))
>>> [0, 1, 2, 1, 2, 3, 2, 3, 4]

Upvotes: 4

kuro
kuro

Reputation: 3226

You can use two loops in list comp -

extra_indices = [index+i for index in indices for i in range(3)]

Upvotes: 6

Related Questions