Trial Division Prime Generator

Question

import math


def generate_primes(limit):
    """Generates and prints a list of primes up to a given limit"""
    prime_list = [2, 3]
    num = 5
    increment_2 = True  # When True increment by 2. When False increment by 4

    while num <= limit:
        is_prime = True
        sqr_lim = math.sqrt(num)
        for x in prime_list:
            # Loop through list of prime numbers and check divisibility
            if x > sqr_lim:
                break
            if num % x == 0:
                is_prime = False
                break

        if is_prime:
            # Add primes to list to only check prime numbers
            prime_list.append(num)

        # Alternate step size between 2 and 4 to skip multiples of 2 and 3
        if increment_2:
            num += 2
            increment_2 = False
        else:
            num += 4
            increment_2 = True

    print(prime_list)
    return prime_list


def main():
    generate_primes(200)


if __name__ == "__main__":
    main()

I'm writing a program to generate a list of prime numbers to a given limit through the use of trial division. I was wondering if there is a way to further optimize this code. Currently it: skips multiples of 2 and 3; stores primes in a list and the list is used to check for primality; and checks divisibility up until the square root of the number it's analyzing.

I would also appreciate it if someone could critique me on my coding style such as comments, white space, and organization because I know its pretty rough around the edges. Thanks for the help.

Answer 1

The function generate_primes() shouldn’t print the list it generated. The caller should print the returned value, if desired.

---

Calling generate_primes(2) will unexpectedly return [2, 3], as will generate_primes(1) or generate_primes(-500).

---

The increment_2 flag is awkward. When True, increment by 2 seems ok, but when False the increment is ... what? 0? 1?

Instead of an increment_2 flag, you could use an increment amount, starting with 2 ...

increment = 2

... and then toggle that value between 2 & 4:

num += increment

# Alternate step size between 2 and 4
increment = 6 - increment

No more if ... else ..., reducing the complexity of the function slightly.

---

Stretch goal:

After generate_primes(200), if a call is made to generate_primes(500), you will recalculate all of the primes beneath 200. Perhaps there is a way to cache the computed values for future generate calls?

Answer 2

You could use the else feature of loops, which gets executed iff no break statement was executed. This gets rid of your flag variable:

while num <= limit:
    sqr_lim = math.sqrt(num)
    for x in prime_list:
        # Loop through list of prime numbers and check divisibility
        if x > sqr_lim:
            break
        if num % x == 0:
            break
    else:
        # Add primes to list to only check prime numbers
        prime_list.append(num)

However, this would mean that the condition that breaks the loop because you have reached the square limit would avoid this from being triggered.

To avoid this, you could have a look at the itertools module, which provides the takewhile function:

for p in takewhile(lambda p: p <= sqr_lim, prime_list):
    if num % p == 0:
        break
else:
    prime_list.append(num)

Your variable increment also boils down to the observation that all prime numbers larger than 3 are of the form 6k + 1 or 6k + 5 with k = 0, 1, .... This is for the same reason you used the variable increment, all numbers of the form 6k + 0, 2, 4 are divisible by two and all numbers of the form 6k + 3 by three.

You could use this with the itertools recipe roundrobin to generate all candidate numbers.

from itertools import cycle, islice

def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Recipe credited to George Sakkis
    num_active = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Remove the iterator we just exhausted from the cycle.
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))

candidates = roundrobin(range(7, limit, 6),  # 6k + 1, k = 1, 2, ...
                        range(5, limit, 6))  # 6k + 5, k = 0, 1, ...
for num in candidates:
    ... 

However, if you want to get all primes up to some large(ish) number, sieves are really hard to beat. The easiest one to implement is the Sieve of Eratosthenes, which can for example be done like this:

def prime_sieve(limit):
    prime = [True] * limit
    prime[0] = prime[1] = False

    for i, is_prime in enumerate(prime):
        if is_prime:
            yield i
            for n in range(i * i, limit, i):
                prime[n] = False

If you need even more speed, you could take a look at the Sieve of Atkin. However, using the above implementation of the Sieve of Eratosthenes, I can generate all primes up to 100.000.000 in less than 30 seconds, which is usually fast enough.