pythonarraysnumpy

Why numpy.array so slow when editing the data inside?


I wrote an algorithm that uses a long list. Since numpy.array should perform better in dealing with long data, I wrote another version using numpy.array.

list version:

import math
def PrimeTable(n:int) -> list[int]:
   nb2m1 = n//2 - 1
   l = [True]*nb2m1
   for i in range(1,(math.isqrt(n)+1)//2):
      if l[i-1]:
         for j in range(i*3,nb2m1,i*2+1):
            l[j] = False
   return [2] + [i for i,v in zip(range(3,n,2), l) if v]

array version:

import math, numpy
def PrimeTable2(n:int) -> list[int]:
   nb2m1 = n//2 - 1
   l = numpy.full(nb2m1, True)
   for i in range(1,(math.isqrt(n)+1)//2):
      if l[i-1]:
         for j in range(i*3,nb2m1,i*2+1):
            l[j] = False
   return [2] + [i for i,v in zip(range(3,n,2), l) if v]

It turns out that numpy.array version is 1x slower than list version. The graph:graphwhere y0 is PrimeTable and y1 is PrimeTable2

Why is numpy.array so slow, and how can I improve it?


Solution

  • Numpy is optimized for vectorized operations, where large chunks of data are processed in bulk. Instead of updating element-by-element(l[j] = False) we can use slicing to update a range of values at once. Also reducing the number of python loops should make the code more effective.

    I've further optimised the code by using a single boolean array. The sieve tracks only odd numbers directly. This eliminates the need for multiplication and division steps for odd indices during iteration. After the sieve is complete, the prime numbers are directly reconstructed from the indices, avoiding intermediate computations.

    Try out this updated code. This should work better than the list version of code:

    import numpy as np
    
    def PrimeTable2_optimized_v2(n: int) -> list[int]:
        if n < 2:
            return []
        if n == 2:
            return [2]
    
        # Only consider odd numbers; even numbers > 2 are not prime
        sieve = np.ones(n // 2, dtype=bool)
        limit = int(n**0.5) + 1
    
        for i in range(1, limit // 2 + 1):
            if sieve[i]:
                start = 2 * i * (i + 1) 
                sieve[start::2 * i + 1] = False
    
        # Convert sieve to primes
        primes = 2 * np.nonzero(sieve)[0] + 1 
        primes[0] = 2 
        return primes.tolist()