Constraints

• 0 ≤ n ≤ 100000
• Use a sieve-based approach; repeated trial division per candidate is not efficient enough
• Aim for near-linear memory usage and about O(L log log L) time, where L is the sieve limit

Examples

Solution

def solution(n, return_last_only=False):
    import math

    if n < 0:
        raise ValueError("n must be non-negative")
    if n == 0:
        return []
    if n == 1:
        return [2]

    def upper_bound(k):
        if k < 6:
            return 15
        return int(k * (math.log(k) + math.log(math.log(k)))) + 10

    limit = upper_bound(n)

    while True:
        size = (limit - 1) // 2
        sieve = bytearray(b"\x01") * size
        root = math.isqrt(limit)
        max_idx = (root - 3) // 2

        for i in range(max_idx + 1):
            if sieve[i]:
                p = 2 * i + 3
                start = (p * p - 3) // 2
                sieve[start::p] = b"\x00" * (((size - 1 - start) // p) + 1)

        if return_last_only:
            count = 1
            for i, is_prime in enumerate(sieve):
                if is_prime:
                    count += 1
                    if count == n:
                        return [2 * i + 3]
        else:
            primes = [2]
            primes.extend(2 * i + 3 for i, is_prime in enumerate(sieve) if is_prime)
            if len(primes) >= n:
                return primes[:n]

        limit *= 2

Time complexity: O(L log log L), where L is the sieve bound and L is about n(log n + log log n). In practice, the biggest constants come from marking composite multiples in the sieve and scanning the sieve to build the answer.. Space complexity: O(L), using an odd-only sieve so the main array stores about L/2 candidate slots..

Hints

1. For n ≥ 6, the nth prime is less than about n(log n + log log n), which is a good starting bound for your sieve.
2. You can cut memory roughly in half by storing only odd numbers in the sieve.