misc. prime number related functions

prime.py

#!/usr/bin/env python

import pudb
import random

# infinite prime number generator
def gen_prime():
	tbl = dict()
	i = 2
	while True:
		if i in tbl:
			pr = tbl[i]
			for p in pr:
				if (p+i) not in tbl:
					tbl[p+i] = [p]
				else:
					tbl[p+i].append(p)
			del tbl[i]
		else:
			tbl[i**2] = [i]
			yield i
		i += 1

def del_mult(ls,p):
	if ls[p] == 0:
		return

	i = p**2
	while i < len(ls):
		ls[i] = 0
		i += p

# prime number list generator
def primes_to(n):
	#pudb.set_trace()
	pr = [i for i in xrange(0,n+1)]
	pr[0] = 0
	pr[1] = 0

	i = 2
	while i**2 < n:
		del_mult(pr,i)
		i += 1

	return [x for x in pr if x > 0]

def primes_first(n):
	g = gen_prime()
	ls = []
	for _ in xrange(n):
		ls.append(g.next())
	return ls

def to_binary(n):
	r = []
	while (n > 0):
		r.append(n % 2)
		n = n / 2
	return r

def complex_test(a, n):
	"""
	complex_test(a, n) -> bool complex_tests whether n is complex.

	Returns:
		- True, if n is complex.
		- False, if n is probably prime.
	"""
	b = to_binary(n - 1)
	d = 1
	for i in xrange(len(b) - 1, -1, -1):
		x = d
		d = (d * d) % n
		if d == 1 and x != 1 and x != n - 1:
			return True # Complex
		if b[i] == 1:
			d = (d * a) % n
	if d != 1:
		return True # Complex
	return False # Prime

# http://snippets.dzone.com/posts/show/4200
def miller_rabin(n,s):
	"""
		miller_rabin(n, s = 1000) -> bool Checks whether n is prime or not

		Returns:
			- True, if n is probably prime.
			- False, if n is complex.
	"""
	for j in xrange(1, s + 1):
		a = random.randint(1, n - 1)
		if (complex_test(a, n)):
			return False # n is complex
	return True # n is prime

# primality test
def is_prime(n,k=20):
	if n < 2:
		return False
	if n==2 or n==3:
		return True
	if n%2==0 or n%3==0:
		return False
	return miller_rabin(n,k)

def is_compisite(n,k=20):
	return not is_prime(n,k)

# naive solution
def find_divisor_naive(n):
	r = [1]
	if n == 1:
		return r

	if pr.is_prime(n):
		r.append(n)
		return r

	for i in xrange(2,n/2+1):
		if n%i == 0:
			r.append(i)

	r.append(n)
	return r

# not working, look into prime factorization
def find_divisor(n):
	pudb.set_trace()
	if n == 1:
		return set([1])

	if pr.is_prime(n):
		return set([1,n])

	ans = set([1,n])
	for i in xrange(2,n/2+1):
		if n%i == 0:
			a = find_div(n/i)
			b = find_div(i)
			#print "a=",a
			#print "b=",b
			ans.union(a).union(b)
	return ans

if __name__ == "__main__":
	#for x in xrange(550):
		#if is_prime(x):
			#print x,
	#for prime in gen_prime():
		#print prime
	print