给定质数求其序号的Python代码超时问题求助

Your current implementation has several key inefficiencies that lead to excessive runtime, especially for larger primes:

1. Inefficient Primality Check: For each number num, you check divisibility up to num/2 instead of the square root of num. This wastes time because any factor larger than sqrt(num) would have a corresponding factor smaller than sqrt(num).
2. Redundant Divisibility Check: You check if n % num == 0 for every num up to n, which is an O(n) operation. Instead, you can verify if n is prime once using an efficient method.
3. Suboptimal Prime Generation: Generating primes by checking each number individually with a slow primality test leads to an overall O(n²) time complexity, which is way too slow for even moderately large primes.

Optimized Solution #

Here's a revised version of your code that fixes these issues, using efficient primality testing and the Sieve of Eratosthenes for fast prime counting:

def kthPrime(self, n):
    # Helper function to check if a number is prime efficiently
    def is_prime(num):
        if num <= 1:
            return False
        if num <= 3:
            return True
        # Quick checks to eliminate even numbers and multiples of 5
        if num % 2 == 0 or num % 5 == 0:
            return False
        # Check divisors up to sqrt(num), stepping by 2 (only odd numbers)
        i = 3
        while i * i <= num:
            if num % i == 0:
                return False
            i += 2
        return True

    # First, verify if n is actually a prime
    if not is_prime(n):
        return 0

    # Sieve of Eratosthenes to count primes up to n efficiently
    def count_primes_sieve(num):
        if num < 2:
            return 0
        sieve = [True] * (num + 1)
        sieve[0] = sieve[1] = False
        for i in range(2, int(num ** 0.5) + 1):
            if sieve[i]:
                # Mark all multiples of i as non-prime
                sieve[i*i : num+1 : i] = [False] * len(sieve[i*i : num+1 : i])
        return sum(sieve)

    # Return the count of primes up to n (which is the position of n)
    return count_primes_sieve(n)

Key Optimizations Explained #

1. Efficient Primality Check:

   • We first eliminate small cases and even numbers/multiples of 5 with quick checks to avoid unnecessary work.
   • We only check odd divisors up to the square root of num, which reduces the number of iterations from O(num) to O(sqrt(num)).

2. Sieve of Eratosthenes:

   • This algorithm generates all primes up to n in O(n log log n) time, which is exponentially faster than your original linear approach.
   • It works by marking multiples of each prime as non-prime, allowing us to count all primes in one pass.

3. Early Prime Verification:

   • We first confirm n is a prime using the efficient helper function. If it's not, we immediately return 0, avoiding unnecessary computation.

Additional Notes #

• For extremely large primes (10^8 or higher), the sieve may use too much memory. In that case, you could switch to a segmented sieve or use a prime-counting algorithm like Meissel-Lehmer, but those are more complex. For most practical purposes, the sieve approach will be more than fast enough.
• The quick checks for even numbers and multiples of 5 help eliminate non-primes early, saving time in both the primality test and sieve.

内容的提问来源于stack exchange，提问作者C. Wei