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.
def isprime(n):
divs = range(2, int(n ** 0.5) + 1)
return all(n % d for d in divs)
%time isprime(1_000_003)
A common optimization is to skip even numbers.
def isprime(n):
divs = range(3, int(n ** 0.5) + 1, 2)
return n == 2 or all(n % d for d in divs)
%time isprime(1_000_003)
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.
import itertools
def isprime(n):
divs = range(5, int(n ** 0.5) + 1, 2)
return n in (2, 3) or all(n % d for d in itertools.chain(divs[::3], divs[1::3]))
%time isprime(1_000_003)
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.
def primes():
yield 2
ints = itertools.count(3, 2)
while True:
prime = next(ints)
yield prime
ints = (n for n in ints if n % prime)
list(itertools.islice(primes(), 10))
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.
def primes():
yield 2
ints = itertools.count(3, 2)
while True:
prime = next(ints)
yield prime
ints = filter(prime.__rmod__, ints)
%time next(itertools.islice(primes(), 1000, None))
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.
def primes():
yield 2
primes = []
for n in itertools.count(3, 2):
if all(n % p for p in itertools.takewhile(int(n ** 0.5).__ge__, primes)):
primes.append(n)
yield n
%time next(itertools.islice(primes(), 1000, None))
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.
class Primes:
def __init__(self):
self.ints = itertools.count(3, 2)
self.cache = [2]
def __iter__(self):
yield from self.cache
for n in self.ints:
if n in self:
self.cache.append(n)
yield n
def __contains__(self, n):
return all(n % p for p in itertools.takewhile(int(n ** 0.5).__ge__, self))
primes = Primes()
%time next(itertools.islice(primes, 1000, None))
%time 1_000_003 in primes
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.
def primes():
multiples = {}
for n in itertools.count(2):
prime = multiples.pop(n, 0)
if not prime:
prime = n
yield n
key = n + prime
while key in multiples:
key += prime
multiples[key] = prime
%time next(itertools.islice(primes(), 1000, None))
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.
def primes():
yield 2
multiples = {}
for n in itertools.count(3, 2):
step = multiples.pop(n, 0)
if step: # composite
key = n + step
while key in multiples:
key += step
multiples[key] = step
else: # prime
multiples[n ** 2] = n * 2
yield n
%time next(itertools.islice(primes(), 1000, None))
And finally let's add back the caching. Yielding a clean interface, an efficient implementation for all use cases, and still relatively simple.
class Primes:
def __init__(self):
self.ints = itertools.count(3, 2)
self.cache = [2]
self.multiples = {}
def __iter__(self):
yield from self.cache
for n in self.ints:
step = self.multiples.pop(n, 0)
if step: # composite
key = n + step
while key in self.multiples:
key += step
self.multiples[key] = step
else: # prime
self.multiples[n ** 2] = n * 2
self.cache.append(n)
yield n
def __contains__(self, n):
return all(n % p for p in itertools.takewhile(int(n ** 0.5).__ge__, self))
primes = Primes()
%time 1_000_003 in primes
%time 1_000_003 in primes
from Planet Python
via read more
No comments:
Post a Comment