Implement scalable prime generator
Company: Pinterest
Role: Data Scientist
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Onsite
Write a function first_n_primes(n) that returns the first n prime numbers in ascending order.
Constraints:
- 1 ≤ n ≤ 100,000.
- Aim for O(n log log n) expected time and near-linear memory.
- Avoid repeated trial division per candidate; use an efficient sieve and a tight upper bound estimate for the nth prime.
- Handle n=0 and n=1 edge cases cleanly.
- Include unit tests for n ∈ {0,1,5,10,100,100000} and verify the last prime for n=100000 equals the known value.
- Describe the complexity and key constants that dominate runtime for large n.
Quick Answer: This question evaluates a candidate's ability to implement scalable prime generation and related algorithmic optimization, including algorithm selection, time and memory complexity reasoning, edge-case handling, and verification via unit tests.
Write a function `solution(n, return_last_only=False)` that generates prime numbers efficiently using a sieve. By default, it should return the first `n` prime numbers in ascending order. To keep very large tests compact, if `return_last_only` is `True`, return a one-element list containing only the `n`th prime instead of the full list. For `n = 0`, return an empty list. Your implementation should avoid repeated trial division and should use a tight upper-bound estimate for the `n`th prime before sieving. For large inputs, the runtime is dominated by crossing off composite numbers in the sieve and, when returning the full answer, by materializing the output list.
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
Input: (0, False)
Expected Output: []
Explanation: There are no primes to return when n is 0.
Input: (1, False)
Expected Output: [2]
Explanation: The first prime number is 2.
Hints
- 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.
- You can cut memory roughly in half by storing only odd numbers in the sieve.