Cfrtjryk

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?

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.

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.

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