Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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

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.

like image 377
InQβ Avatar asked Oct 15 '25 12:10

InQβ


1 Answers

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()

enter image description here

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.

like image 71
MSeifert Avatar answered Oct 17 '25 01:10

MSeifert