Python: Prime Numbers

Here’s a simple code for checking if a given number is prime:

def isPrime(n):
 
	for i in range(2, n):
		if n%i == 0:
			return False
	return True
 
print(isPrime(2)) # true
print(isPrime(3)) # true
print(isPrime(11)) # true
print(isPrime(8)) # false

For the input 2 you may be wondering how it returns True. This is because we never enter the for loop for that input. Notice the range becomes 2,2. Since the range goes from “first argument” to “second argument – 1”, it will go from 2 to 1 which doesn’t make sense. Which is why we never enter the for loop.

It has a worst case complexity of O(n) because the for loop will run from 2 to n in worst case.

Can we do better than that? Yes, we can. We only need to check till square root of the given number. Using this logic we re-write the code as follows:

import math
 
def isPrime2(n):
 
	n2 = int(math.sqrt(n))
 
	for i in range(2,n2+1):
		if n%i == 0:
			return False
	return True
 
print(isPrime2(2)) # true
print(isPrime2(3)) # true
print(isPrime2(8)) # false

You may be wondering why we changed the range to (2,n2+1). This is because if we were to use (2,n), the input 8 wouldn’t return the correct result.

The square root of 8 is 2.82 but it would be returned as 2. The for loop would not be entered in this case and the output would be True which is incorrect for input 8. This is fixed by making it n2+1.

The worst case time complexity here would be O(sqrt(n)).

Leave a comment