improve performance of list creation

Viewed 418

How can I improve significantly the speed of the following code? Can mapping, numpy, matrix operations be efficiently used and/or something else to omit the for loop?

import time

def func(x):
    if x%2 == 0:
        return 'even'
    else:
        return 'odd'

starttime = time.time()

MAX=1000000

y=list(range(MAX))

for n in range(MAX):
    y[n]=[n,n**2,func(n)]

print('That took {} seconds'.format(time.time() - starttime))

The following replacement does not improve the speed:

import numpy as np
r = np.array(range(MAX))
str = ['odd', 'even']
result = np.array([r, r ** 2, list(map(lambda x: str[x % 2], r))])
y = result.T
7 Answers

I think you can do it this way, the idea is to use as many numpy built-in functions as possible

%%timeit
y = np.arange(MAX)
y_2 = y**2
y_str = np.where(y%2==0,'even','odd')

res = np.rec.fromarrays((y,y_2,y_str), names=('y', 'y_2', 'y_str'))

#
# Some examples for working with the record array
res[3]
# (3, 9, 'odd')
res[:3]
# rec.array([(0, 0, 'even'), (1, 1, 'odd'), (2, 4, 'even')],
#           dtype=[('y', '<i8'), ('y_2', '<i8'), ('y_str', '<U4')])
res['y_str'][:7]
# array(['even', 'odd', 'even', 'odd', 'even', 'odd', 'even'], dtype='<U4')
res.y_2[:7]
# array([ 0,  1,  4,  9, 16, 25, 36])

I have ran several tests, and it is significantly faster.

For large arrays of the same type, numpy is the way to go. But numpy where is slow, so if you just want 'odd' and 'even', you can use np.tile or something like it:

MAX = 1000000

%%timeit
y = np.arange(MAX)
ystr = np.where(y%2==0,'even','odd')
#  14.9 ms ± 61.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
temp = np.array(['even', 'odd'])
ystr = np.tile(temp, MAX//2)
# 4.1 ms ± 112 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

So tile is about 3-4x faster.

If you want something more complex, I'd still try to avoid where if speed is important. There's almost always a way because the where logic is usually simple so it's easy to take the logical expression that was inside the where and write it as an expression between numpy arrays. (Also, to be sure, using numpy and where will be much faster than pure Python lists, it's just usually slow relative to other numpy options.)

The others are fairly obvious:

y = np.arange(MAX) 
y2 = y**2

Personally, I'd just stick these together in a list,

result = [y, y2, ystr]

Putting this all together (using tile), I get:

# 6.82 ms ± 84.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

The short answer would be you can't.

Let's explore this a bit through examples.

# Original code
y=list(range(MAX))
for n in range(MAX):
    y[n]=[n,n**2,func(n)]

# That took 0.86 seconds

This is the result on my machine so we have a baseline for comparison. Let us make that a single line and shave off some time.

y = [[n, n ** 2, func(n)] for n in range(MAX)]

# That took 0.74 seconds

We are creating a list of lists and the Python interpreter needs to allocate an empty list MAX times.

In case you don't need to change the number of elements after initialization it might be better to use tuples instead of lists.

y = [(n, n ** 2, func(n)) for n in range(MAX)]

# That took 0.43 seconds

This is twice as fast as the original method. Let's now assume that we can optimize even more by using some special library and then we just need to parse the result to populate the list. To simulate this we can pickle the list to a binary format and then measure the time it takes to load it.

import pickle

b = pickle.dumps([(n, n ** 2, func(n)) for n in range(MAX)])

starttime = time.time()
y = pickle.loads(b)
print('That took {:.2f} seconds'.format(time.time() - starttime))

# That took 0.23 seconds

This is probably close to what is possible to achieve without coding anything in a lower level language like C and creating Python objects from that language.

Alternative approach

If there is no requirement to create exactly the same object as in the original example and if it is enough that we can read y[10] or y[100:1000] we can do something completely different.

class LazyList():
    def __init__(self, size):
        self.size = size

    def __getitem__(self, key):
        if isinstance(key, slice):
            r = range(self.size)[key]
            return [(n, n ** 2, func(n)) for n in r]

        return (key, key ** 2, func(key))

starttime = time.time()
y = LazyList(MAX)

print('That took {:.6f} seconds'.format(time.time() - starttime))

# That took 0.000005 seconds

This is multiple orders of magnitude faster. Of course, this is not a list and the results of the computation are not in memory. We created an object that will in some cases act like a list, but not in other cases (e.g. y[MAX*2] will work, even though it shouldn't). Note that with more work, the object can become even more similar to a list and also use a list as its base class.

If the object we got is converted to a list, the process spends the time that was saved by the alternative approach and the result is the same as in one of the previous examples.

y = y[:]

# That took 0.43 seconds

The longer answer is that it depends on the type of the result that is expected.

See below (it takes ~ 1 sec on my mac)

import time

MAX = 1000000

starttime = time.time()

y = [[n, n ** 2, 'even' if n % 2 == 0 else 'odd'] for n in range(MAX)]

print('That took {} seconds'.format(time.time() - starttime))

I would recommend you to run it with multiple processes in parallel:

import time
import multiprocessing


def func(x):
    return [x, x ** 2, "even" if x % 2 == 0 else "odd"]


if __name__ == '__main__':
    starttime = time.time()
    MAX = 1000000
    pool = multiprocessing.Pool(10)
    y = pool.map(func, range(MAX))
    print('That took {} seconds'.format(time.time() - starttime))

Try to tune the number of processes to get the optimal value for your environment. On mine, it took ~0.8 secs with 20 processes while your original snippet took ~1.1 secs.

Here is a numpy approach:

import time
import numpy as np
starttime = time.time()

r = np.arange(MAX)
res = [r, r ** 2, np.where(r % 2, 'odd', 'even')]

print('That took {:.4} seconds'.format(time.time() - starttime))
# That took 0.05125 seconds  || original function took 1.5s

As @Divakar pointed out, how to move on from here depends on what end result you want. One option would be to have an object array with mixed types:

res = np.array(res, dtype=object).T
print('That took {:.4} seconds'.format(time.time() - starttime))
# That took 0.1863 seconds

res[17]
# array([17, 289, 'odd'], dtype=object)
res[18] + res[17] 
# array([35, 613, 'evenodd'], dtype=object)  # add for int and str

Unfortunately it is quite expansive to combine the 3 different arrays. It is still way faster than using loops but depending on your next steps you could maybe make further improvements.

On my computer:

  • the original loop took about 1.01 seconds
  • NumPy solution took 10.3 ms
  • Numba solution took 4.25 ms
from numba import njit
import numpy as np

def f(n_max = 1_000_000):
    y = x ** 2
    z = x % 2
    return y, z

@njit
def g(x):
    y = x ** 2
    z = x % 2
    return y, z

n_max = 1_000_000
x = np.arange(n_max, dtype=int)

NumPy:

%%timeit
y, z = f(x)
10.3 ms ± 296 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

And Numba:

y, z = g(x) # don't time first run, which does compile AND execute

%%timeit
y, z = g(x)
4.25 ms ± 85.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Related