3

This one is clear:

a = ['a', 'b', 'c']
x = [v for v in a]

But I am not sure how to do this:

sep = 5
# ??? y = [v + sep for v in a]
print y # expected ['a', 5, 'b', 5, 'c', 5]

How can I write a list comprehension with multiple elements per source element?

I am not interested in optimizations of this code: please do not refer to the [:] operator or join method or something on those lines. My code needs a list comprehension. The only alternative I have at the moment is a 4 lines for loop, which is inconvenient:

y = []
for v in a:
    y.append(v)
    y.append(sep)
blueFast
  • 41,341
  • 63
  • 198
  • 344

6 Answers6

6

Using nested list comprehension:

>>> a = ['a','b','c']
>>> [item  for x in a for item in (x, 5)]
['a', 5, 'b', 5, 'c', 5]
falsetru
  • 357,413
  • 63
  • 732
  • 636
5

You can build a list of tuples first, and then flatten the resulting list using itertools.chain.from_iterable:

>>> sep = 5
>>> a = ['a', 'b', 'c']
>>> 
>>> import itertools
>>> list(itertools.chain.from_iterable((elem, sep) for elem in a))
['a', 5, 'b', 5, 'c', 5]
Rohit Jain
  • 209,639
  • 45
  • 409
  • 525
  • As @RohitJain was faster, let me add that a list comprehension will give you at most one element in your result for each element of the list that you traverse. You want to have two. I believe this is only possible with something along the lines that Rohit suggests. – nickie Oct 07 '13 at 16:47
  • @nickie. I don't understand your comment. – Rohit Jain Oct 07 '13 at 16:50
  • The comment is meant as an explanation why this cannot be done just with one list comprehension traversing `a` (i.e., without the extra flattening). If you understand it and can write it better, please do. If you think there's no hope, I'll erase it. – nickie Oct 07 '13 at 16:54
  • @nickie. Oh! No leave it. No issues. – Rohit Jain Oct 07 '13 at 16:55
1

Simply use the function sum() to flatten the list of list as [[v,sep] for v in a] will produce [['a', 5], ['b', 5], ['c', 5]]

a = ['a', 'b', 'c']
x = [v for v in a]
sep = 5
y = sum([[v,sep] for v in a],[])
print y
#['a', 5, 'b', 5, 'c', 5]
K DawG
  • 13,287
  • 9
  • 35
  • 66
1

This is a perfect problem to show how powerful Python's itertool package is

repeat: Creates an endless (or upto a limit) iterable of an object

izip: Transposes the iterables.

chain Unwraps an iterable

>>> from itertools import izip, chain, repeat
>>> a = ['a', 'b', 'c']
>>> list(chain.from_iterable(izip(a, repeat(5))))
['a', 5, 'b', 5, 'c', 5]

And if you have a knack for micro optimization, you may be interested to know, this is faster than list comprehension

>>> stmt1 = """
list(chain.from_iterable((elem, 5) for elem in a))
"""
>>> stmt2 = """
list(chain.from_iterable(izip(a, repeat(5))))
"""
>>> timeit.timeit(stmt=stmt1, setup="from __main__ import chain, a", number=100000)
1.3136643729533688
>>> timeit.timeit(stmt=stmt2, setup="from __main__ import chain, izip, repeat, a", number=100000)
0.8959859753707633
>>>
Abhijit
  • 62,056
  • 18
  • 131
  • 204
0
result = [5 if i % 2 else a[i/2] for i in range(2*len(a))]

(Not that you should put code like this into production)

YXD
  • 31,741
  • 15
  • 75
  • 115
0
>>> a = ['a', 'b', 'c']
>>> [item for sublist in zip(a, [5]*len(a)) for item in sublist]
['a', 5, 'b', 5, 'c', 5]

To break it down:

>>> [5]*len(a)
[5, 5, 5]

>>> zip(a, [5]*len(a))
[('a', 5), ('b', 5), ('c', 5)]

The list of tuples is then flattened using the following idiom:

[item for sublist in l for item in sublist]

Using timeit, this approach was 15% faster than the fastest itertools method

>>> stmt3 = """
[item for sublist in zip(a, [5]*len(a)) for item in sublist]
"""
>>> timeit.timeit(stmt=stmt3, setup="a = ['a', 'b', 'c']", number=100000)
dansalmo
  • 11,506
  • 5
  • 58
  • 53