Thursday, 25 April 2013

Prime Numbers

A few months ago, during a lecture on Haskell and functional programming, we were introduced to the Sieve of Eratosthenes, an elegant way to generate prime numbers. We were shown a neat implementation in Haskell that used just a few lines; I wanted to see if I could replicate that neatness/laziness in Python.

from itertools import takewhile, chain, count, islice
from collections import deque

knownPrimes = [2] # Numbers are tested by counting upwards from here. This also acts as a cache.

def storeAndReturn(p):
    knownPrimes.append(p)
    return p

def isNotDivisibleByKnownPrimes(n):
    return all((n % p) != 0 for p in takewhile(lambda p: p**2 <= n, knownPrimes))

def primeFilter(sequence): return (storeAndReturn(i) for i in sequence if isNotDivisibleByKnownPrimes(i))

def allPrimes(): return chain(knownPrimes, primeFilter(count(knownPrimes[-1] + 1)))

def primesUpTo(limit): return takewhile(lambda p: p < limit, allPrimes())

def firstNPrimes(n): return islice(allPrimes(), None, n, None)

def nthPrime(n): return deque(firstNPrimes(n), maxlen=1).pop() # deque construction is apparently the
# fastest way to consume an iterator (partly because it's raw C)

if __name__ == '__main__': # Module Test
    print "Primes below 10:"
    print '\n'.join(map(str, primesUpTo(10)))
    
    print "First 10 primes:"
    print '\n'.join(map(str, firstNPrimes(10)))
    
    print "10,000th prime:"
    print nthPrime(10000)

I succeeded, and even added some more interesting ways to use the allPrimes() generator (although a few of the one-liners should probably be broken-out to aid readability).

Primes below 10:
2
3
5
7
First 10 primes:
2
3
5
7
11
13
17
19
23
29
10,000th prime:
104729

Neat, huh?

I love infinite generators - it's flat-out cool to be able to create an object describing every prime number using just a few lines.

No comments:

Post a Comment

About Me

Melbourne, Victoria, Australia
I'm hard at work on an MSc in Physics at the University of Melbourne, my research topic is 'building nano-scale neural networks'. In my (limited) spare time I tinker with 3D printing, electronics and programming.