import itertools
def generate_n_primes(n):
"""
Generate n prime numbers using a modified Sieve of Eratosthenes.
The algorithm keeps track of a list of primes found so far,
and a corresponding list of 'multiples', where multiples[i] is a multiple of primes[i],
(multiples[i] is initially set to be primes[i]**2, see the optimisations section below).
The main loop iterates over every integer k until enough primes have been found,
with the following steps:
- For each prime found so far
- While the corresponding multiple is smaller than k, increase it by steps of the prime
- If the multiple is now the same as k, then k is divisible by the prime -
hence k is composite, ignore it.
- If, for EVERY prime, the multiple is greater than k, then k isn't divisible by any
of the primes found so far. Hence we can add it to the prime list and multiple list!
There are a few optimisations that can be done:
- We can insert 2 into primes at the start, and only iterate over every odd k from there on
- When we're increasing the multiple, we can now increase by 2*prime instead of 1*prime,
so that we skip over even numbers, since we are now only considering odd k
- When we find a prime p, we add it to the prime and multiple list. However, we can instead add
its square to the multiple list, since for any number between p and p**2, if it's
divisible by p then it must be divisible by another prime k < p
(i.e. it will be caught by an earlier prime in the list)
"""
# Insert 2 into primes/multiples
primes = [2]
multiples = [4]
# Iterate over odd numbers starting at 3
for k in itertools.count(3, 2):
# If we've found enough primes, return!
if len(primes) >= n:
return primes
# For each prime found so far
for i in range(len(primes)):
# Increase its corresponding multiple in steps of 2*prime until it's >= k
while multiples[i] < k:
multiples[i] += 2 * primes[i]
# If its corresponding multiple == k then k is divisible by the prime
if multiples[i] == k:
break
else:
# If k wasn't divisible by any prime, add it to the primes/multiples list
primes.append(k)
multiples.append(k ** 2)
return primes
Some might find the docstring as well as comments too much - I find the comments help relate the code to the docstring. Open to suggestions!
