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 ?

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))