Finding nth Prime using Python and Sieve of Eratosthenes

Question

I'm currently working through the Project Euler Problems using HackerRank to evaluate the code and I'm stuck on the 7th Problem.

"""
    By listing the first six prime numbers: 2, 3, 5, 7, 11, and 13, we can see that the 6th prime is 13.

    What is the 10 001st prime number?

"""
import pytest

def find_primes(limit):
    nums = [True] * (limit+1)
    nums[0] = nums[1] = False

    for (i, is_prime) in enumerate(nums):
        if is_prime:
            yield i
            for n in range(i*i, limit+1, i):
                nums[n] = False

def find_n_prime(n):
    for i in range(n, (n*n)+1, n):
        primes = list(find_primes(i))
        if len(primes) >= n:
            return primes[n-1]

def test_find_n_primes():
    assert find_n_prime(2) == 3
    assert find_n_prime(3) == 5

    assert find_n_prime(10) == 29
    assert find_n_prime(15) == 47
    assert find_n_prime(81) == 419

    assert find_n_prime(941) == 7417
    assert find_n_prime(1000) == 7919

    assert find_n_prime(10001) == 104743

The code only completes the first test case on hackerRank and fails for #2 test case and times out for the rest.

How can I improve the Code ?

Answer 1

Upper bound for p_n

There is a known upper bound for the n-th prime.

It means that you don't need any loop inside find_n_prime, and you don't need to check if len(primes) >= n either.

import pytest
from math import log, ceil

def find_primes(limit):
    nums = [True] * (limit + 1)
    nums[0] = nums[1] = False

    for (i, is_prime) in enumerate(nums):
        if is_prime:
            yield i
            for n in range(i * i, limit + 1, i):
                nums[n] = False

def upper_bound_for_p_n(n):
    if n < 6:
        return 100
    return ceil(n * (log(n) + log(log(n))))

def find_n_prime(n):
    primes = list(find_primes(upper_bound_for_p_n(n)))
    return primes[n - 1]

It calculates the 100000th prime in less than 230ms on my computer, compared to 1.5s for your code.

itertools.islice

Another possible optimization would be to use itertools.isclice to get the n-th prime out of the find_primes generator, without converting it to a list.

from itertools import islice

def find_n_prime(n):
    primes = find_primes(upper_bound_for_p_n(n))
    return next(islice(primes, n - 1, n))