Thursday, August 12, 2010

Simple prime test in Python

A simple way to test if a number is prime is to check if the number is divisible by any other number besides one and itself. This can take a lot of time, but for small numbers, it might be ok.
def is_prime(n):
   # assumes n > 0
   if n < 2: return True
   for d in range(2,n):
       if n % d == 0: return False 
   return True
This simple function can be sped up by not checking as many numbers. We don't need to test all divisors up to n-1. The largest possible divisor is n/2, so we could stop there. But, even better, we can show that if there are any divisors, there must be divisor less than the sqrt(n). If p*q = n, then either p or q must be less than or equal to sqrt(n). Otherwise, p*q would be bigger than n. So, we can just loop up to sqrt(n) because we only need to find a divisor of n, not every divisor. Here is the updated code:
import math
def is_prime2(n):
    if n < 2: return False
    for d in xrange(2, int(math.sqrt(n)) + 1):
        if n % d == 0: return False
    return True

>>> [x for x in xrange(40) if is_prime2(x)]
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]

Finally, we could use the same code to find a divisor of a number. In fact, it finds the smallest divisor. Some of the more advanced primality tests do not determine any of the factors of the number.
def is_prime3(n):
    """returns the smallest divisor of n.  
       returns 1 if the number is prime, or less than 2"""
    if n < 2: return 1
    for d in xrange(2, int(math.sqrt(n)) + 1):
        if n % d == 0: return d
    return 1

2 comments:

njzk2 said...

i would suggest to loop on

xrange(3, (int) math.sqrt(n) + 1, 2)

just to try only half as many possibilities
{off course, it also means testing for '2' first)

Dave said...

Hi Simon,

You are right. I planned to use this simple function to verify the answers from more advanced functions, but I got a little distracted and never did it.