# Generate nth prime number in Python

Today I had an interview, where I was asked to solve this problem:

Generate nth prime number. Given a signature below, write python logic to generate the nth prime number:

def nth_prime_number(n):
# n = 1 => return 2
# n = 4 => return 7
# n = 10 => return 29


I wrote this code, but couldn't get through:

def nth_prime_number(n):
if n==1:
return 2
count = 1
num = 3
while(count <= n):
if is_prime(num):
count +=1
if count == n:
return num
num +=2 #optimization

def is_prime(num):
factor = 2
while (factor < num):
if num%factor == 0:
return False
factor +=1
return True


The overall feedback was that the code quality could be improved a lot, and that I should be more optimal in my approach. How can I improve this code?

• Was the output incorrect, or was there some other reason? – Jamal Mar 26 '17 at 18:34
• The feedback was that the general code quality is not good, that's all, and that I need to be more optimal in my apporach. – Harsha Mar 26 '17 at 18:41

Your is_prime() function checks if num is a multiple of any number below it. This means that it checks if it is a multiple of 2, 4, 6, 8, 10, etc. We know that if it isn't a multiple of 2, it won't be a multiple of 4, etc. This goes for all other numbers, if it isn't a multiple of 3, it won't be a multiple of 27 (3x3x3).

What you need to do is check if num is a multiple of any prime number before it.

def nth_prime_number(n):
# initial prime number list
prime_list = [2]
# first number to test if prime
num = 3
# keep generating primes until we get to the nth one
while len(prime_list) < n:

# check if num is divisible by any prime before it
for p in prime_list:
# if there is no remainder dividing the number
# then the number is not a prime
if num % p == 0:
# break to stop testing more numbers, we know it's not a prime
break

# if it is a prime, then add it to the list
# after a for loop, else runs if the "break" command has not been given
else:
# append to prime list
prime_list.append(num)

# same optimization you had, don't check even numbers
num += 2

# return the last prime number generated
return prime_list[-1]


I'm sure someone else here will come up with an even more efficient solution, but this'll get you started.

• This seems to go further than the "sieve" (see below answers) by stating that we only need to check whether a number is divisible by primes that are less than it, and not by any other numbers. Can you provide a reference to vouch for the correctness of this method? – Stephen Jul 9 '18 at 21:57
• I'm not going to spend ages searching for a reference, but this post on Math seems to put it nicely: If a number isn't prime (and isn't 1), it's divisible by a number other than itself and 1, which means it's divisible by the prime factors of that number. – DarkMatterMatt Jul 10 '18 at 0:47
• Surprisingly this is slower than this answer which divides by all numbers < sqrt(num), see benchmark. However you can combine both solutions to get an even better one. – Socowi Nov 26 '18 at 15:08

You only need to loop to the square root of the number, because after that you'll just repeat the same numbers again. For example if you were testing for 100, after 10 you will find 20, but you already tested it while you were testing the 5, 100/5 = 20 and 100/20 = 5, if 5 didn't divide 100 so 20 won't and vice versa, so 100/a = b tests for the divisibility by a and b, only at 10 which is sqrt(100) will a and b be repeated again but in reversed order, (e.g you a=5, b=20 and a=20, b=5). More info here

So the code will look like that

def nth_prime_number(n):
if n==1:
return 2
count = 1
num = 1
while(count < n):
num +=2 #optimization
if is_prime(num):
count +=1
return num

def is_prime(num):
factor = 2
while (factor * factor <= num):
if num % factor == 0:
return False
factor +=1
return True


But overall this naive method consumes a lot of time, it's O(sqrt(n) * n) where n is the total number of integers tested, I suggest you learn more about Sieve of Eratosthenes, it's an algorithm for finding all primes less than n in O(n * log(log(n))) which is alot faster than the naive one.

• This is a good answer except for the last bit about the Sieve of Eratosthenes. The Sieve would only be useful if you had an known upper bound and you wanted to know all of the primes below it. – ljeabmreosn Jun 15 '17 at 2:18
• @ljeabmreosn, you can actually implement an unbounded sieve - I posted one in C++ here for review. – Toby Speight Jul 10 '18 at 8:31
• The code shown has a number of basic bugs in it (failing to return a value from nth_prime_number, for example, and calculating that the first 3 primes are 2,7,9 rather than 2,3,5). – jarmod Dec 6 '18 at 22:41
• @jarmod Wow, I wrote this answer almost 2 yrs ago and no one noticed this major bug. I fixed it. Thanks! – Khaled Hamed Dec 8 '18 at 14:56
• is_prime should use a for loop, not a while loop.
• nth_prime_number should use while True, rather than while count <= n, as you'll never meet that condition.
• #optimization is of no help, how's it an optimization?
• nth_prime_number would be better written as two functions an infinite generator, and a function that picks the nth prime.
• is_prime can be significantly shortened if you use any.

This can get you:

from itertools import count, islice

def is_prime(num):
return any(
num % factor
for factor in range(2, num)
)

def generate_primes():
yield 2
for num in count(3, 2):
if is_prime(num):
yield num

def nth_prime_number(n):
return next(islice(generate_prime(), n, None))


You use the increment by two optimization, but you don't need to check numbers that are greater than $\sqrt{n}$. And you can increment in multiples of six, and picking two numbers to the side of it - $6n−1$ and $6n+1$.

However it'd be best if you used a sieve, such as the Sieve of Eratosthenes. Where you'd base your algorithm off this Math.SE answer.