nth_prime.py

n-th prime (projanmo forum quiz)

#!/usr/bin/env python

from itertools import islice, count

def eratosthenes_improved(n):
    ''' Slightly improved version of Eratosthenes Sieve algorithm.
    This algorithm do not store an array of numbers to rule out composites.
    Instead it uses a dictionary of special composite numbers only, hence saving
    memory. It also rules out even numbers in the loop, making less iteration. '''

    n = int(n)
    # make sure n is an integer
    if n < 1: return None
    # we're looking for n-th prime!
    if n == 1: return 2
    # 2 is the only prime that is even

    sieve = {}
    # here we store our known composite numbers with lowest prime factor
    counter = 1
    for q in islice(count(3), 0, None, 2):
        # we'll only iterate over odd numbers to find a prime
        p = sieve.pop(q, None)
        if p is None:
            # q is a prime, since it is not in known composites
            sieve[q*q] = q
            # q^2 is a composite with q prime factor
            counter += 1
            if counter == n:
                # we found the n-th prime
                break
        else:
            # Now (N*p)+q is also composite
            # let's find the next such composite that is not marked yet
            # and mark it with the lowest prime factor p
            x = p + q
            while x in sieve or not (x&1):
                x += p
            sieve[x] = p

    return q

print eratosthenes_improved(9999) # 104723