How to conditionally put multiple elements at a time in list comprehension?

Question

I have discovered this very close to my question:

Adding two items at a time in a list comprehension

But what if I need to switch between single or double, like:

original = list(range(10))
required = [0,0,1,2,2,3,4,4,5,6,6,7,8,8,9]
attempt1 = sum([[x,x] if x%2 == 0 else [x] for x in original],[])
attempt2 = [i for x in original for i in ([x,x] if x%2 == 0 else [x])]

sum seems slow, and list comprehension is hard to read. Neither of them makes me feel simple and good.

Is there a better way to do it? Or just abandon the one-line way? Or convince me if one of them is really good style.

This boils down to what's the most readable way to un-nest a list. Have a look at `itertools.chain` and `itertools.chain.from_iterable`. — timgeb, Aug 15 '17 at 10:28

MSeifert · Accepted Answer · 2017-08-15T11:52:59.653

4

Personally I would use a generator function as soon as I have more than non-trivial stuff going on in a comprehension (for example 2 fors and one if).

For example in your case you could use (I think it's more readable but that might be subjective):

def double_evens(inp):
    for item in inp:
        if item % 2 == 0:
            yield item
        yield item

Test run:

>>> list(double_evens(range(10)))
[0, 0, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8, 9]

Note that this approach could even be faster (it's 3 times faster than the other solutions in the answers and 2 times faster than your comprehension on my computer). Taking the timing framework from this answer:

from itertools import chain

def coldspeed1(mylist):
    return [y for x in mylist for y in [x] * (2 - x % 2)]

def coldspeed2(mylist):
    return list(chain.from_iterable([x] * (2 - x % 2) for x in mylist))

def double_evens(inp):
    for item in inp:
        if not item % 2:
            yield item
        yield item

def mseifert(inp):
    return list(double_evens(inp))  

def ettanany(my_list):
    new_list = [[i] * 2 if i % 2 == 0 else i for i in my_list]
    res = []
    for i in new_list:
        if isinstance(i, list):
            res.extend(i)
        else:
            res.append(i)
    return res

def no1xsyzy(original):
    return [i for x in original for i in ([x,x] if x%2 == 0 else [x])]

# Timing setup
timings = {coldspeed1: [], coldspeed2: [], mseifert: [], ettanany: [], no1xsyzy: []}
sizes = [2**i for i in range(1, 20, 2)]

# Timing
for size in sizes:
    mylist = list(range(size))
    for func in timings:
        res = %timeit -o func(mylist)
        timings[func].append(res)

# Plotting
%matplotlib notebook

import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure(1)
ax = plt.subplot(111)

baseline = mseifert # choose one function as baseline
for func in timings:
    ax.plot(sizes, 
            [time.best / ref.best for time, ref in zip(timings[func], timings[baseline])], 
            label=str(func.__name__))
#ax.set_yscale('log')
ax.set_xscale('log')
ax.set_xlabel('size')
ax.set_ylabel('time relative to {}'.format(baseline.__name__))
ax.grid(which='both')
ax.legend()
plt.tight_layout()

This graph plots the relative time difference compared to my solution. Note that the x-axis (the sizes) is logarithmic while the y-axis (time difference) isn't.

edited Aug 15 '17 at 11:52

answered Aug 15 '17 at 10:42

MSeifert

145,886
38
333
352

1

While the question says list comprehension in the title and while this is quite far from it, I still have to respect the effort you've put into this answer. Thank you for your service to SO! (+1) – cs95 Aug 15 '17 at 12:18
Hmm, I admit, `Explicit is better than implicit.` and `Readability counts.`, however, it would be too complex to make a single function just for one-time use. Despite that, it should be better than the one-line way. – InQβ Aug 15 '17 at 14:22
BTW, could you also include timing of a variation of your function where different occasions are handled separately? I would like to know how much more time would it cost. – InQβ Aug 15 '17 at 14:24
@no1xsyzy What do you mean by "different occasions"? – MSeifert Aug 15 '17 at 14:29
And how about this coldspeed1 variation: `[x for x in mylist for _ in range(2 - x % 2)]` – InQβ Aug 15 '17 at 14:30
@no1xsyzy Well, technically a comprehension is just a function in disguise so from a programming perspective it isn't much different. And if you think it's more readable and I've shown it's faster than why bother with a "shorter" but otherwise inferior way? It's a judgement call but I've never had a situation that writing a function was "too complex" even for "one-time-use". – MSeifert Aug 15 '17 at 14:30
@no1xsyzy That variation is even slower than the `[y for x in range(10) for y in [x] * (2 - x % 2)]` approach. – MSeifert Aug 15 '17 at 14:33
@MSeifert By "different occasions", I mean yield once/twice in different `if` branches instead of once in branch, once outside. That shouldn't be a big question – InQβ Aug 15 '17 at 14:35
@no1xsyzy Do you mean like [this (link to a github gist)](https://gist.github.com/MSeifert04/f692b641076d92c5aa3d6ef4b06de39f)? In that case the approaches are (almost) equally fast. – MSeifert Aug 15 '17 at 14:39
@MSeifert I would take the advice. Maybe I struggle too much to make the program look like math formulas(functional?). – InQβ Aug 15 '17 at 14:43

cs95 · Answer 2 · 2017-08-15T12:21:37.797

1

You're on the right track. The simplification you can make is through an elimination of and if..else within the list comp by substituting it for a * operation.

In [53]: [y for x in range(10) for y in [x] * (2 - x % 2)]
Out[53]: [0, 0, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8, 9]

This roughly translates to:

res = []
for x in range(10):
    if not (x + 1) % 2:
        res.extend([x])
    else:
        res.extend([x, x])

Alternatively, you could create a list of lists and use itertools.chain.from_iterable to flatten it:

In [54]: import itertools

In [55]: list(itertools.chain.from_iterable([x] * (2 - x % 2) for x in range(10)))
Out[55]: [0, 0, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8, 9]

Appreciate the help from people in the comments for refactoring some of these solutions.

edited Aug 15 '17 at 12:21

answered Aug 15 '17 at 10:28

cs95

379,657
97
704
746

I think `list(itertools.chain.from_iterable([x] if (x + 1) % 2 else [x, x] for x in range(10)))` is a bit more readable (also you don't need to pass a list comprehension to `from_iterable`, just the generator is enough and lazier). – jdehesa Aug 15 '17 at 10:34
@jdehesa Yes... it looks better now. Thank you. Applied the same to first example too. – cs95 Aug 15 '17 at 10:36
Slightly shorter: `[y for x in range(10) for y in [x]*(x % 2 + 1)]` – donkopotamus Aug 15 '17 at 10:41
@donkopotamus Not quite. Had to edit it a bit to get it as OP wanted. – cs95 Aug 15 '17 at 10:45
You could also use `([x,x], [x])[x%2]` (not recommended but it looks funny with all the x's). :) – MSeifert Aug 15 '17 at 10:49
Ah, hadn't seen the mismatch ... then `[x] * (2 - x % 2)` is a bit less unwieldy than `[x] * ((x + 1) % 2 + 1)` – donkopotamus Aug 15 '17 at 10:49

ettanany · Answer 3 · 2017-08-15T10:45:04.037

0

A simple readable solution may look like below:

my_list = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
new_list = [[i] * 2 if i % 2 == 0 else i for i in my_list]
res = []
for i in new_list:
    if isinstance(i, list):
        res.extend(i)
    else:
        res.append(i)

Output:

>>> res
[0, 0, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8, 9]

edited Aug 15 '17 at 10:45

answered Aug 15 '17 at 10:36

ettanany

19,038
9
47
63

your method will nasty if the original list is list of lists, e.g., `[[0],[1],...,[9]]` (explicitly with list comprehension: `[[i] for i in range(10)]`) – InQβ Aug 16 '17 at 05:26

How to conditionally put multiple elements at a time in list comprehension?

3 Answers3