How can I use a list comprehension to extend a list in python?

Question

I'm not experienced in Python, and I often write code that (simplified) looks like this:

accumulationList = []
for x in originalList:
    y = doSomething(x)
    accumulationList.append(y)
return accumulationList

Then after my test passes, I refactor to

return [doSomething(x) for x in originalList]

But suppose it turns out a little different, and my loop looks like this:

accumulationList = []
for x in originalList:
    y = doSomething(x)
    accumulationList.extend(y)
return accumulationList

where the doSomething list returns a list. What is the most Pythonic way to accomplish this? Obviously, the previous list comprehension would give a list of lists.

@chris_rands -- It seems a little odd to mark a question as a dupe 7.5 years after it was asked -- it doesn't serve the purpose of directing the OP to the duplicate, and the questions are not quite similar enough the merging is an option. /shrug — Eric Wilson, Nov 19 '18 at 13:23

score 19 · Accepted Answer · answered Jun 08 '17 at 12:18

19

Much simpler and cleaner with list comprehension:

[y for x in originalList for y in doSomething(x)]

answered Jun 08 '17 at 12:18

Shubham Chaudhary

47,722
9
78
80

1

Great improvement! Back when I asked the question, I wouldn't have though to use a nested list comprehension. – Eric Wilson Jun 08 '17 at 13:52

eumiro · Answer 2 · 2011-05-10T08:48:06.367

5

Do you mean something like this?

accumulationList = []
for x in originalList:
    accumulationList.extend(doSomething(x))
return accumulationList

or shorter code (but not optimal):

return sum((doSomething(x) for x in originalList), [])

or the same:

return sum(map(doSomething, originalList), [])

Thanks to @eyquem for the hint (if using Python 2.x):

import itertools as it

return sum(it.imap(doSomething, originalList), [])

edited May 10 '11 at 08:48

answered May 10 '11 at 08:19

eumiro

207,213
34
299
261

1

``return sum(map(doSomething, originalList)), [])`` better with **imap** if Python 2. Right if Python 3 – eyquem May 10 '11 at 08:28

score 3 · Answer 3 · answered May 10 '11 at 21:26

3

I think the answers involving add or iadd run in quadratic time, which probably isn't good. I'd try:

from itertools import chain
accumulation_list = list(chain.from_iterable(doSomething(x) for x in originalList))

answered May 10 '11 at 21:26

solrize

31
1

score 2 · Answer 4 · answered May 10 '11 at 08:29

2

Python's in-place add operator (+=, available as iadd in operator module) is equivalent of .extend for list. Pair it with reduce to get what you want.

import operator

reduce(operator.iadd, (doSomething(x) for x in originalList)
, accumulation_list)

answered May 10 '11 at 08:29

Imran

87,203
23
98
131

score 0 · Answer 5 · answered Nov 12 '18 at 15:29

Functionally, you can use itertools.chain with map. For an input list L:

res = list(chain.from_iterable(map(doSomething, L)))

If you need an iterator, simply remove the list call. Here's a demo:

def doSomething(value):
    return [value * i for i in range(1, 4)]

def original(originalList):
    accumulationList = []
    for x in originalList:
        y = doSomething(x)
        accumulationList.extend(y)
    return accumulationList

def new(L):
    return list(chain.from_iterable(map(doSomething, L)))

x = [1, 2, 3]

assert original(x) == new(x)

print(new(x))

[1, 2, 3, 2, 4, 6, 3, 6, 9]

score 0 · Answer 6 · answered May 10 '11 at 08:43

I don't think there is special syntax for this case. But you could make the for loop shorter:

accumulationList += doSomething(x)

If you insist, you could use functional programming to flatten the list:

result = reduce(lambda a,b: a+b, [[i,i*2] for i in range(3)])

But I wouldn't call this pythonic, I think it's harder to read than a for loop.

How can I use a list comprehension to extend a list in python?

6 Answers6

Linked

Related