CODEWITHSUNDEEP
pythonarraysperformancepython-3.xnumpy
Jano
Jano

Reputation: 94

Array initialized with numpy is much slower in a loop with if statement

I have written a prime generator. To improve the performance i tried to use numpy for a big array.

[True for x in range(prime_max)] takes 4.5 seconds at all with prime_max = 10000000, but 0.728 to initialize

np.ones(prime_max, dtype=bool) takes 18.9 seconds at all with prime_max = 10000000 , but only 0.004 to initialize

is this because of if number_list[i] == True: ? Is there a faster way?

def prime(prime_max):
    prime_list = []
    number_list = [True for x in range(prime_max)]  
    #number_list = np.ones(prime_max, dtype=bool)   
    for i in range(2, prime_max):
        if number_list[i] == True:
            prime_list.append(i)
            for j in range(2*i, prime_max, i):
                number_list[j] = False
    return prime_list

I know there are faster ways to generate primes, but I am learning python and first I want to understand this problem and optimize this code.

Upvotes: 1

Views: 177

Answers (1)

user2357112
user2357112

Reputation: 281538

NumPy is optimized for whole-array operations, not single-element accesses. Single-element accesses are quite slow in NumPy, since every access needs to resolve stuff like dimensions and dtypes dynamically, and element retrieval needs to create a pretty heavyweight wrapper object on demand every time you retrieve an element.

Efficient NumPy code avoids Python-level loops over arrays. For example, replacing

for j in range(2*i, prime_max, i):
    number_list[j] = False

with

number_list[2*i::i] = False

would probably speed up your NumPy-based code a great deal.

Here's how I'd write this with NumPy:

def primes(prime_max):
    prime_flags = numpy.ones(prime_max, dtype=bool)
    prime_flags[0] = prime_flags[1] = False
    for i in range(2, prime_max):
        if prime_flags[i]:
            prime_flags[2*i::i] = False
    return prime_flags.nonzero()[0]

We can't get rid of the outer Python loop easily, but we can make some improvements. (I'd also stop the loop at sqrt(prime_max), but I tried to focus on improvements in NumPy usage rather than algorithmic improvements.)

Upvotes: 4

Related Questions