Pythonicity: Primes

An old interview challenge is to generate prime numbers or check if a number is prime. No advanced mathematics needed, just variants on the Sieve of Eratosthenes. Starting with a basic prime checker.

A common optimization is to skip even numbers.

Brief digression on that optimization. There’s nothing special about removing multiples of 2; removing multiples is the whole point. The step scalar could instead be thought of as a cycle: itertools.accumulate(itertools.repeat(2)). So removing multiples of 3 would remove every third step: itertools.accumulate(itertools.cycle([2, 4])).

Or the equivalent could be done with slicing.

The catch is the cycles grow exponentially with diminishing returns on each successive number.

Onto prime generation, while keeping the odds-only optimization. Typically it’s requested to generate the first N primes, or up to some value. But that’s easily generalized with itertools.islice and itertools.takewhile. A more Pythonic approach is an unbounded generator.

Elegant, but doesn’t work. The problem is the scoping of prime, which is being used in the generator expression but also modified in the loop. Instead it can be replaced with a filter on a partially bound function, but unfortunately functools.partial only binds left arguments and rmod is needed here. One alternative is to use bound methods as a first-class function, even dunder methods.

Elegant, but slow and could overflow the stack. A more traditional approach would use the same checking logic as isprime, but also cache the primes so as to not duplicate divisors.

Onto interface design. The primes are being stored anyway, so it would be nice if they were re-iterable. A generator can be written as a class with __iter__ and __next__, but an under-appreciated feature is that __iter__ itself can be a generator. And now that it’s a class, isprime can be expressed as in while also benefiting from the cache.

There’s a hybrid approach though, that’s faster and nearly as simple as the above sieves. Instead of doing repeated divisions, keep track of each found prime along with the next multiple that it would eliminate. The inner loop is then optimized because it only needs to account for collisions.

Now to add back the odds-only optimization, the step scalar needs to be double the prime number. Another way to reduce collisions is to recognize that each new prime is irrelevant until its square value is reached.

And finally let’s add back the caching. Yielding a clean interface, an efficient implementation for all use cases, and still relatively simple.