Speed up bitstring/bit operations in Python?

There are a couple of small optimizations for your version. By reversing the roles of True and False, you can change "if flags[i] is False:" to "if flags[i]:". And the starting value for the second range statement can be i*i instead of i*3. Your original version takes 0.166 seconds on my system. With those changes, the version below takes 0.156 seconds on my system.

def prime_numbers(limit=1000000):    '''Prime number generator. Yields the series    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...    using Sieve of Eratosthenes.    '''    yield 2    sub_limit = int(limit**0.5)    flags = [True, True] + [False] * (limit - 2)    # Step through all the odd numbers    for i in range(3, limit, 2):        if flags[i]:            continue        yield i        # Exclude further multiples of the current prime number        if i <= sub_limit:            for j in range(i*i, limit, i<<1):                flags[j] = True

This doesn't help your memory issue, though.

Moving into the world of C extensions, I used the development version of gmpy. (Disclaimer: I'm one of the maintainers.) The development version is called gmpy2 and supports mutable integers called xmpz. Using gmpy2 and the following code, I have a running time of 0.140 seconds. Running time for a limit of 1,000,000,000 is 158 seconds.

import gmpy2def prime_numbers(limit=1000000):    '''Prime number generator. Yields the series    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...    using Sieve of Eratosthenes.    '''    yield 2    sub_limit = int(limit**0.5)    # Actual number is 2*bit_position + 1.    oddnums = gmpy2.xmpz(1)    current = 0    while True:        current += 1        current = oddnums.bit_scan0(current)        prime = 2 * current + 1        if prime > limit:            break        yield prime        # Exclude further multiples of the current prime number        if prime <= sub_limit:            for j in range(2*current*(current+1), limit>>1, prime):                oddnums.bit_set(j)

Pushing optimizations, and sacrificing clarity, I get running times of 0.107 and 123 seconds with the following code:

import gmpy2def prime_numbers(limit=1000000):    '''Prime number generator. Yields the series    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...    using Sieve of Eratosthenes.    '''    yield 2    sub_limit = int(limit**0.5)    # Actual number is 2*bit_position + 1.    oddnums = gmpy2.xmpz(1)    f_set = oddnums.bit_set    f_scan0 = oddnums.bit_scan0    current = 0    while True:        current += 1        current = f_scan0(current)        prime = 2 * current + 1        if prime > limit:            break        yield prime        # Exclude further multiples of the current prime number        if prime <= sub_limit:            list(map(f_set,range(2*current*(current+1), limit>>1, prime)))

Edit: Based on this exercise, I modified gmpy2 to accept xmpz.bit_set(iterator). Using the following code, the run time for all primes less 1,000,000,000 is 56 seconds for Python 2.7 and 74 seconds for Python 3.2. (As noted in the comments, xrange is faster than range.)

import gmpy2try:    range = xrangeexcept NameError:    passdef prime_numbers(limit=1000000):    '''Prime number generator. Yields the series    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...    using Sieve of Eratosthenes.    '''    yield 2    sub_limit = int(limit**0.5)    oddnums = gmpy2.xmpz(1)    f_scan0 = oddnums.bit_scan0    current = 0    while True:        current += 1        current = f_scan0(current)        prime = 2 * current + 1        if prime > limit:            break        yield prime        if prime <= sub_limit:            oddnums.bit_set(iter(range(2*current*(current+1), limit>>1, prime)))

Edit #2: One more try! I modified gmpy2 to accept xmpz.bit_set(slice). Using the following code, the run time for all primes less 1,000,000,000 is about 40 seconds for both Python 2.7 and Python 3.2.

from __future__ import print_functionimport timeimport gmpy2def prime_numbers(limit=1000000):    '''Prime number generator. Yields the series    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...    using Sieve of Eratosthenes.    '''    yield 2    sub_limit = int(limit**0.5)    flags = gmpy2.xmpz(1)    # pre-allocate the total length    flags.bit_set((limit>>1)+1)    f_scan0 = flags.bit_scan0    current = 0    while True:        current += 1        current = f_scan0(current)        prime = 2 * current + 1        if prime > limit:            break        yield prime        if prime <= sub_limit:            flags.bit_set(slice(2*current*(current+1), limit>>1, prime))start = time.time()result = list(prime_numbers(1000000000))print(time.time() - start)

Edit #3: I've updated gmpy2 to properly support slicing at the bit level of an xmpz. No change in performance but a much nice API. I have done a little tweaking and I've got the time down to about 37 seconds. (See Edit #4 to changes in gmpy2 2.0.0b1.)

from __future__ import print_functionimport timeimport gmpy2def prime_numbers(limit=1000000):    '''Prime number generator. Yields the series    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...    using Sieve of Eratosthenes.    '''    sub_limit = int(limit**0.5)    flags = gmpy2.xmpz(1)    flags[(limit>>1)+1] = True    f_scan0 = flags.bit_scan0    current = 0    prime = 2    while prime <= sub_limit:        yield prime        current += 1        current = f_scan0(current)        prime = 2 * current + 1        flags[2*current*(current+1):limit>>1:prime] = True    while prime <= limit:        yield prime        current += 1        current = f_scan0(current)        prime = 2 * current + 1start = time.time()result = list(prime_numbers(1000000000))print(time.time() - start)

Edit #4: I made some changes in gmpy2 2.0.0b1 that break the previous example. gmpy2 no longer treats True as a special value that provides an infinite source of 1-bits. -1 should be used instead.

from __future__ import print_functionimport timeimport gmpy2def prime_numbers(limit=1000000):    '''Prime number generator. Yields the series    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...    using Sieve of Eratosthenes.    '''    sub_limit = int(limit**0.5)    flags = gmpy2.xmpz(1)    flags[(limit>>1)+1] = 1    f_scan0 = flags.bit_scan0    current = 0    prime = 2    while prime <= sub_limit:        yield prime        current += 1        current = f_scan0(current)        prime = 2 * current + 1        flags[2*current*(current+1):limit>>1:prime] = -1    while prime <= limit:        yield prime        current += 1        current = f_scan0(current)        prime = 2 * current + 1start = time.time()result = list(prime_numbers(1000000000))print(time.time() - start)

Edit #5: I've made some enhancements to gmpy2 2.0.0b2. You can now iterate over all the bits that are either set or clear. Running time has improved by ~30%.

from __future__ import print_functionimport timeimport gmpy2def sieve(limit=1000000):    '''Returns a generator that yields the prime numbers up to limit.'''    # Increment by 1 to account for the fact that slices do not include    # the last index value but we do want to include the last value for    # calculating a list of primes.    sieve_limit = gmpy2.isqrt(limit) + 1    limit += 1    # Mark bit positions 0 and 1 as not prime.    bitmap = gmpy2.xmpz(3)    # Process 2 separately. This allows us to use p+p for the step size    # when sieving the remaining primes.    bitmap[4 : limit : 2] = -1    # Sieve the remaining primes.    for p in bitmap.iter_clear(3, sieve_limit):        bitmap[p*p : limit : p+p] = -1    return bitmap.iter_clear(2, limit)if __name__ == "__main__":    start = time.time()    result = list(sieve(1000000000))    print(time.time() - start)    print(len(result))

OK, so this is my second answer, but as speed is of the essence I thought that I had to mention the bitarray module - even though it's bitstring's nemesis :). It's ideally suited to this case as not only is it a C extension (and so faster than pure Python has a hope of being), but it also supports slice assignments. It's not yet available for Python 3 though.

I haven't even tried to optimise this, I just rewrote the bitstring version. On my machine I get 0.16 seconds for primes under a million.

For a billion, it runs perfectly well and completes in 2 minutes 31 seconds.

import bitarraydef prime_bitarray(limit=1000000):    yield 2    flags = bitarray.bitarray(limit)    flags.setall(False)    sub_limit = int(limit**0.5)    for i in xrange(3, limit, 2):        if not flags[i]:            yield i            if i <= sub_limit:                flags[3*i:limit:i*2] = True

Okay, here's a (near complete) comprehensive benchmarking I've done tonight to see which code runs the fastest. Hopefully someone will find this list useful. I omitted anything that takes more than 30 seconds to complete on my machine.

I would like to thank everyone that put in an input. I've gained a lot of insight from your efforts, and I hope you have too.

My machine: AMD ZM-86, 2.40 Ghz Dual-Core, with 4GB of RAM. This is a HP Touchsmart Tx2 laptop. Note that while I may have linked to a pastebin, I benchmarked all of the following on my own machine.

I will add the gmpy2 benchmark once I am able to build it.

All of the benchmarks are tested in Python 2.6 x86

Returning a list of prime numbers n up to 1,000,000: (Using Python generators)

Sebastian's numpy generator version (updated) - 121 ms @

Mark's Sieve + Wheel - 154 ms

Robert's version with slicing - 159 ms

My improved version with slicing - 205 ms

Numpy generator with enumerate - 249 ms @

Mark's Basic Sieve - 317 ms

casevh's improvement on my original solution - 343 ms

My modified numpy generator solution - 407 ms

My original method in the question - 409 ms

Bitarray Solution - 414 ms @

Pure Python with bytearray - 1394 ms @

Scott's BitString solution - 6659 ms @

'@' means this method is capable of generating up to n < 1,000,000,000 on my machine setup.

In addition, if you don't need the generator and just want the whole list at once:

numpy solution from RosettaCode - 32 ms @

(The numpy solution is also capable of generating up to 1 billion, which took 61.6259 seconds. I suspect the memory was swapped once, hence the double time.)