Checking if a number is a prime number in Python

Question

I have written the following code, which should check if the entered number is a prime number or not, but there is an issue i couldn't get through:

def main():
n = input("Please enter a number:")
is_prime(n)

def is_prime(a):
    x = True 
    for i in (2, a):
            while x:
               if a%i == 0:
                   x = False
               else:
                   x = True


    if x:
        print "prime"
    else:
        print "not prime"

main()

If the entered number is not a prime number, it displays "not prime", as it is supposed to, but if the number is a prime number, it doesn't display anything. Could you please help me with it?

Answer 1

Here is my take on the problem:

from math import sqrt
from itertools import count, islice

def is_prime(n):
    return n > 1 and all(n % i for i in islice(count(2), int(sqrt(n)-1)))

This is a really simple and concise algorithm, and therefore it is not meant to be anything near the fastest or the most optimal primality check algorithm. It has a time complexity of O(sqrt(n)). Head over here to learn more about primality tests done right and their history.

---

Explanation

I'm gonna give you some insides about that almost esoteric single line of code that will check for prime numbers:

• First of all, using range() in Python 2 is really a bad idea, because it will create a list of numbers, which uses a lot of memory. Using xrange() is better, because it creates a generator, which only needs to memorize the initial arguments you provide, and generates every number on-the-fly. If you're using Python 3, range() has been converted to a generator by default. By the way, this is still not the best solution: trying to call xrange(n) for some n such that n > 231-1 (which is the maximum value for a C long) raises OverflowError. Therefore the best way to create a range generator is to use itertools:

   xrange(2147483647+1) # OverflowError
  
   from itertools import count, islice
  
   count(1)                        # Count from 1 to infinity with step=+1
   islice(count(1), 2147483648)    # Count from 1 to 2^31 with step=+1
   islice(count(1, 3), 2147483648) # Count from 1 to 3*2^31 with step=+3
  

• You do not actually need to go all the way up to n if you want to check if n is a prime number. You can dramatically reduce the tests and only check from 2 to √(n) (square root of n). Here's an example:

  Let's find all the divisors of n = 100, and list them in a table:

   2  x  50 = 100
   4  x  25 = 100
   5  x  20 = 100
  10  x  10 = 100 <-- sqrt(100)
  20  x  5  = 100     
  25  x  4  = 100
  50  x  2  = 100
  

  You will easily notice that, after the square root of n, all the divisors we find were actually already found. For example 20 was already found doing 100/5. The square root of a number is the exact mid-point where the divisors we found begin being duplicated. Therefore, to check if a number is prime, you'll only need to check from 2 to sqrt(n).

• Why sqrt(n)-1 then, and not just sqrt(n)? That's just because the second argument provided to itertools.islice() is the number of iterations to execute. islice(count(a), b) stops after b iterations. That's the reason why:

   for number in islice(count(10), 2):
       print number,
  
   # Will print: 10 11
  
   for number in islice(count(1, 3), 10):
       print number,
  
   # Will print: 1 4 7 10 13 16 19 22 25 28
  

• The function all(...) is the same of the following:

   def all(iterable):
       for element in iterable:
           if not element:
               return False
       return True
  

  It literally checks for all the elements in the iterable, returning False when any of them evaluates to False (which for an integer means only if it's zero). Why do we use it then? First of all, we don't need to use an additional index variable (like we would do using a loop), other than that: just for concision, there's no real need of it, but it looks way less bulky to work with only a single line of code instead of several nested lines.

Extended version

I'm including an "unpacked" version of the isPrime() function, to make it easier to understand and read:

from math import sqrt
from itertools import count, islice

def is_prime(n):
    if n < 2:
        return False

    for number in islice(count(2), int(sqrt(n) - 1)):
        if n % number == 0:
            return False

    return True