Generate nth prime number in Python

Question

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?

Answer 1

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.

Answer 2

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.

Answer 3

• 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.

Answer 4

You can use Wilson's theorem to seriously optimize the is_prime() function. It is also a good idea to turn it into a lambda function if you also want shorter code.

factorial = lambda a:a and a*x(a-1)or 1
is_prime = lambda b:factorial(b-1)%b==b-1

You should then nest a while loop in a for loop to get the nth prime.