Prime sieve generator

Question

I wrote a simple prime sieve in Python. But it ended up looking... ugly. I included a few of the most basic optimizations. When posting this question SE suggested a bunch of questions on the same topic, which lead me to some improvements. But few were in generator form.

It's the optimizations' fault.
Unfortunately I can't continue an outer loop from an inner loop, like I can in JavaScript. I very much disagree with Guido here, as it makes the inner loop clunky.

from itertools import count

def sieve():
    primes = [2]
    for candidate in count(start=3, step=2):
        cont = False
        n = 0
        while primes[n]**2 <= candidate:  # You only need to check up to the square root of a number.
            if candidate % primes[n] == 0:
                cont = True  # outer
                break
            n = n + 1
        if cont:
            cont = False
            continue
        yield primes[-1]
        primes.append(candidate)

Can we make this more concise without changing the basic logic? For example this is extremely concise, but it doesn't have some of the optimizations my code does.

For fun, I wrote the same logic in Javascript. This looks cleaner due to being able to continue an outer loop. But the lack of negative indices is a step backwards.

function* sieve() {
    let primes = [2]
    counter: for (let candidate = 3;; candidate+=2) {
        for (let n = 1; primes[n]**2 <= candidate; n++)
            if (candidate % primes[n] == 0)
                continue counter
        yield primes[primes.length - 1]
        primes.push(candidate)
    }
}

Answer 1

1. We can move the while loop into another function.

   def is_prime(candidate, primes):
       n = 0
       while primes[n]**2 <= candidate:
           if candidate % primes[n] == 0:
               return False
           n += 1
       return True
   
   def sieve():
       primes = [2]
       for candidate in count(start=3, step=2):
           if not is_prime(candidate, primes):
               continue
           yield primes[-1]
           primes.append(candidate)
   

2. We can use itertools.takewhile to expose the while loop as an iterable.

   def is_prime(candidate, primes):
       for prime in itertools.takewhile(lambda p: p**2 <= candidate, primes):
           if candidate % prime == 0:
               return False
       return True
   

3. We can use any to make is_prime easier to read.

   If for any of the values, candidate % prime == 0, are true the result is true. If none of them are then it is false. Since we want it to be the other way around we can just use not.

   def is_prime(candidate, primes):
       return not any(
           candidate % prime == 0
           for prime in itertools.takewhile(lambda p: p**2 <= candidate, primes)
       )
   

4. We can move is_prime back into the first function.

   def sieve():
       primes = [2]
       for candidate in count(start=3, step=2):
           if not not any(
               candidate % prime == 0
               for prime in itertools.takewhile(lambda p: p**2 <= candidate, primes)
           ):
               continue
           yield primes[-1]
           primes.append(candidate)
   

5. We can swap the if to cancel the continue and a not.

6. We can swap the any and == 0, with not all.
7. We can cancel the double not.

def sieve():
    primes = [2]
    for candidate in count(start=3, step=2):
        if all(
            candidate % prime
            for prime in itertools.takewhile(lambda p: p**2 <= candidate, primes)
        ):
            yield primes[-1]
            primes.append(candidate)

At the expense of readability you can get the following trade offs.

1. There is a chance that using int(candidate ** 0.5).__ge__ is faster than the lambda.

Answer 2

This outputs the same as your function:

import array
from itertools import count


def sieve():
    primes = array.array("L", [2])  # L: unsigned long int
    for candidate in count(start=3, step=2):
        n = 0
        while primes[n] ** 2 <= candidate:
            if candidate % primes[n] == 0:
                break
            n += 1
        else:  # nobreak
            yield primes[-1]
            primes.append(candidate)


n = 0
primes = sieve()

while True:
    prime = next(primes)
    print(prime)
    n += 1

I did not touch the algorithm itself because I am unfamiliar, but Python (tested on 3.8.2, but also runs on 2.7.18, as you seem to need given your tag python-2.x) has the while/else construct that can help you with your control flow here.

A break break, break <label> or similar thinkable constructs have been proposed, as you linked as well, but rejected. A possible remedy is to extract sub-routines into functions and to use their return statements for control flow/multi-level breaking. In your case, that is not needed.

The while/else construct is unfamiliar to many. In fact, it is so alien that Guido would not implement it again nowadays. For now, it is best to think of the else as nobreak: the while loop finished normally (its condition evaluated to False) and exited. Since no break occurred, the else block is executed.

In the above case, if break is hit in the if block, the else is skipped: no yield occurs, and since after the else block, there is no code left, a continue for the outer for loop is implied and not explicitly necessary, since there is nothing else to do anyway.

---

A frequent example is in the form of for/else (which works like the while/else) when looking for a hit, like in your case:

for file in files:
    if file == file_looked_for:
        file.do_something()
        break
else:  # nobreak
    # code to handle file not found

So while, like you, I have been tripped up by the lack of advanced break functionalities, I agree with Guido. Apart from the unfortunate naming of else in the while construct (more discussion here), it can do the job just fine. Trying to break through multiple levels is an occasion to rethink the implementation.

---

Other observations:

• primes[n] has to call __getitem__, which happens twice. This is constant-time, but the following will probably provide a speed-up (but requires the "walrus" operator from Python 3.8):

    while (
        (prime := primes[n]) ** 2 <= candidate
    ):  # You only need to check up to the square root of a number.
        if candidate % prime == 0:
  

  primes no longer has to be looked up, the simple, local prime suffices for the second call.

• n = n + 1 can be n += 1. However, this is possibly slower.

• Finally, the built-in array.array can be much faster than a list. This is because lists can hold arbitrary objects, whereas an array has to be initiliazed for a specific one (here: L for long int, which will last a while). As a positive consequence, the array can then be optimized accordingly.

  Note how array.array, in this case, is a 1:1 drop-in for the previous list. As such, the code did not have to change. As such, it is also easy for you to revert the change if unwanted.

Answer 3

Three main things vex me about your code:

• the just generated prime isn't returned,
• the boolean test a the end of the while loop
• repeated prime "square and test" for candidacy.

Not returning generated prime

Your code generates 3, and yields 2, then it generates 5 and yields 3, then it generates 7 and yields 5, then generates 11 and yields 7, and so on.

This happens because you half treat 2 as a special case. You initialize the primes array with it. But to return it, you use yield primes[-1] just like every other prime.

If you treated it completely as a special case, and yield it right off the hop, you could yield candidate at the end of the loop, thus returning the prime you just computed.

def sieve():
    primes = [2]
    yield 2

    for candidate in count(start=3, step=2):
        ...
        yield candidate
        primes.append(candidate)

Unnecessary boolean test at end of while

A while loop is often used for searching. If the value is found, the while loop is escaped via a break statement. If the while loop completes without ever breaking, the condition searched for was never found, and something different needs to happen. In Python, this is the while ... else statement:

def sieve():
    primes = [2]
    yield 2
    for candidate in count(start=3, step=2):
        n = 0
        while primes[n]**2 <= candidate:  # Only check up to the square root of number.
            if candidate % primes[n] == 0:
                break
            n = n + 1
        else:
            yield candidate
            primes.append(candidate)

Repeated prime "square and test" for candidacy.

How often is the primes[n]**2 <= candidate done?

If candidate is just over 10,000, and is prime, then we will be squaring all primes less than 100, and testing that they are less than candidate. Then, we do the same thing for candidate + 2, and the results will be the same. No prime number less than 100, squared, will ever be greater than candidate once candidate exceeds 10,000 ... so this is all busy work, repeating the same test over and over.

What you need is to partition your primes list into two parts: primes less or equal to the square-root of candidate, and primes greater the square-root of candidate.

You can do this in several ways. The smallest change would be to keep track of a count of "small" primes. As candidate gets larger by 2, you would only need add at most one more prime into the "small" primes bucket:

def sieve():
    primes = [2]
    yield 2

    small_primes = 0
    for candidate in count(start=3, step=2):

        if primes[small_primes] ** 2 <= candidate:
            small_primes += 1

        for n in range(small_primes):
             if candidate % primes[n] == 0:
                break
        else:
            yield candidate
            primes.append(candidate)

Now how often is primes[small_primes] ** 2 <= candidate being done? Once per candidate! This has got to be an improvement. Also, all n = 0 and n = n + 1 code has been absorbed into for n in range(small_primes), and having Python do this work is faster than coding it ourselves.

Other improvements

Odd numbers

Why are we test-dividing all of our candidates by primes[0] == 2? By design, they are all odd, and can never be evenly divided by 2.

        for n in range(1, small_primes):    # Skip divide-by-2 tests

All

As mentioned by Peilonrayz, Python has an any() function, though I think all() is more appropriate here.

def sieve():
    primes = [2]
    yield 2

    small_primes = 0
    for candidate in count(start=3, step=2):

        if primes[small_primes] ** 2 <= candidate:
            small_primes += 1

        if all(candidate % primes[n] != 0 for n in range(1, small_primes)):
            yield candidate
            primes.append(candidate)

Maintain separate lists

Instead of small_primes being a count of the number of primes less than the square-root of the candidate, what if it actually was a list of the small prime numbers? And instead of adding prime candidates to that list, we add to a large_primes list? Then we could move primes from the large_primes to the small_primes as the square-root of the candidate increases.

Optimizations:

• deque for large_primes
• Omit 2 from the small_primes list,
• Cache the large_prime[0] ** 2 value, to avoid repeatedly squaring the same quantity.

Resulting code:

from itertools import count
from collections import deque

def sieve():
    yield 2
    yield 3

    small_primes = []
    large_primes = deque((3,))
    next_prime_squared = large_primes[0] ** 2

    for candidate in count(start=5, step=2):

        if candidate >= next_prime_squared:
            small_primes.append(large_primes.popleft())
            next_prime_squared = large_primes[0] ** 2

        if all(candidate % prime != 0 for prime in small_primes):
            yield candidate
            large_primes.append(candidate)

Time Comparisons

Time (in seconds) for generating 100 to 100,000 primes:

Timing code:

import array
from timeit import timeit
from itertools import count, takewhile
from collections import deque
import matplotlib.pyplot as plt

def martixy():
    primes = [2]
    for candidate in count(start=3, step=2):
        cont = False
        n = 0
        while primes[n]**2 <= candidate:  # You only need to check up to the square root of a number.
            if candidate % primes[n] == 0:
                cont = True  # outer
                break
            n = n + 1
        if cont:
            cont = False
            continue
        yield primes[-1]
        primes.append(candidate)

def alex_povel():
    primes = array.array("L", [2])  # L: unsigned long int
    for candidate in count(start=3, step=2):
        n = 0
        while primes[n] ** 2 <= candidate:
            if candidate % primes[n] == 0:
                break
            n += 1
        else:  # nobreak
            yield primes[-1]
            primes.append(candidate)

def peilonrayz():
    primes = [2]
    for candidate in count(start=3, step=2):
        if all(
            candidate % prime
            for prime in takewhile(lambda p: p**2 <= candidate, primes)
        ):
            yield primes[-1]
            primes.append(candidate)

def ajneufeld():
    yield 2
    yield 3

    small_primes = []
    large_primes = deque((3,))
    next_prime_squared = large_primes[0] ** 2

    for candidate in count(start=5, step=2):

        if candidate >= next_prime_squared:
            small_primes.append(large_primes.popleft())
            next_prime_squared = large_primes[0] ** 2

        if all(candidate % prime != 0 for prime in small_primes):
            yield candidate
            large_primes.append(candidate)

def test(candidate, limit):
    sieve = candidate()
    for _ in range(limit):
        next(sieve)

if __name__ == '__main__':
    candidates = (martixy, alex_povel, peilonrayz, ajneufeld)
    limits = [int(10 ** (power * 0.25)) for power in range(8, 21)]

    fig, ax = plt.subplots()

    for candidate in candidates:
        print("Testing", candidate.__name__)
        times = [ timeit(lambda: test(candidate, limit), number=1) for limit in limits ]
        ax.plot(limits, times, '-+', label=candidate.__name__)

    ax.legend()
    plt.show()