prime.py

import heapq
import sys

if sys.version_info >= (3, 0):
    xrange = range


def is_prime_trial_division(n):
    """Prime checking by trial division, checking only odds up to sqrt(n).

    >>> is_prime_trial_division(2)
    True
    >>> is_prime_trial_division(3)
    True
    >>> is_prime_trial_division(4)
    False
    >>> is_prime_trial_division(16129)
    False
    >>> is_prime_trial_division(16381)
    True
    """
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    d = 3
    while d * d <= n:
        if n % d == 0:
            return False
        d += 2
    return True


def is_prime_sieve(n):
    """Prime checking by building a Sieve of Eratosthenes.

    Ideally, you would build the sieve up to the maximum number to be tested,
    and re-use it for testing a range of numbers.

    >>> is_prime_sieve(2)
    True
    >>> is_prime_sieve(3)
    True
    >>> is_prime_sieve(4)
    False
    >>> is_prime_sieve(31)
    True
    >>> is_prime_sieve(16129)
    False
    >>> is_prime_sieve(16381)
    True
    """
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    sieve = [True] * (n // 2 + 1)
    for i in xrange(3, n // 3 + 1, 2):
        if not sieve[i // 2]:
            continue
        for j in xrange(3 * i, n + 1, 2 * i):
            sieve[j // 2] = False
    return sieve[n // 2]


def is_prime_sieve_iterative(n):
    """Prime checking with a modified sieve.  Instead of filling out the whole
    sieve at once, only keep track of the next multiple of each confirmed prime.

    >>> is_prime_sieve_iterative(2)
    True
    >>> is_prime_sieve_iterative(3)
    True
    >>> is_prime_sieve_iterative(4)
    False
    >>> is_prime_sieve_iterative(31)
    True
    >>> is_prime_sieve_iterative(16129)
    False
    >>> is_prime_sieve_iterative(16381)
    True
    """
    if n < 2:
        return False
    if n in (2, 3):
        return True
    if n % 2 == 0:
        return False
    sieve = [(3, 9)]
    for i in xrange(5, n, 2):
        is_prime = True
        while True:
            j, prime = sieve[0]
            if i < j:
                break
            is_prime = False
            heapq.heapreplace(sieve, (j + 2 * prime, prime))
        if is_prime:
            heapq.heappush(sieve, (3 * i, i))
    return n < sieve[0][0]


if __name__ == '__main__':
    import doctest
    doctest.testmod()