Search code examples
pythonpython-3.xlist-comprehension

How to save repeated computation in list comprehension in Python?


In the following Python code:

keyboards = [3, 1]
drivers = [5, 2, 8]
upper_limit = 10
sums = [k + d for k in keyboards for d in drivers if (k + d) <= upper_limit]

I'd like to store the result of k+d in the list comprehension so that it can be referred in the list comprehension. Is it possible in Python3?

I know that we can do the following:

sums = []
for k in keyboards:
    for d in drivers:
        s = k + d
        if s <= upper_limit:
            sums.append(s)

But I wish to avoid the side effect operation of append.


Solution

  • In 3.8 and above, you can use the walrus operator for this:

    sums = [s for k in keyboards for d in drivers if (s := k + d) <= upper_limit]
    

    The performance advantage appears to be slight for this example:

    $ python -m timeit "[s for k in range(1000) for d in range(1000) if (s := k + d) <= 1000]"
    5 loops, best of 5: 72.5 msec per loop
    $ python -m timeit "[k + d for k in range(1000) for d in range(1000) if k + d <= 1000]"
    5 loops, best of 5: 75.6 msec per loop
    

    The k + d computation would, presumably, have to be a lot more complex to show a major benefit.