Sometimes we want to create collections of elements that are very expensive to calculate. The first option is to create a list and wait until all the elements are calculated before we use it. Although this works, it is not very efficient. To make it a bit more efficient, modern languages provide a way to create custom iterators so each element is calculated only when needed (this is also called lazy initialization). Also, iterators allows us to create infinite collections!
Python has, in my opinion, one of the most succint and elegant ways to declare iterators: generators.
Without further ado, let's try to create an infinite list of prime numbers.
The first thing we need is a way to detect if a number is prime:
from math import sqrt
def is_prime(n):
if (n <= 1):
return False
if (n == 2):
return True
if (n % 2 == 0):
return False
i = 3
while i <= sqrt(n):
if n % i == 0:
return False
i = i + 2
return True
Now, using our is_prime function we can do:
def prime_generator():
n = 1
while True:
n += 1
if is_prime(n):
yield n
And that's it! Just call the function and get elements from it:
generator = prime_generator()
for i in range(10):
print(next(generator))
Or create a list of the first N prime numbers:
from itertools import islice
array = [x for x in islice(prime_generator(), 10)]
As you could see, the iterator definition is one of the shortest and simplest among all languages.
Top comments (12)
Since the time complexity of
is_primeisO(sqrt(n)):Worth a mention that for prime numbers up to
N, overall, using this method is not as time efficient as sieve methods that eagerly pre-calculate all the primes up toN(this is because sieves use previously found primes to eliminate other composite numbers, instead of performing an exhaustive search for each number).Hi,
Amazing post. It's a really great example to understand how to use generators in Python!
I feel the code above can be improved a bit further, by starting the check for prime numbers from a more suitable number, hence the primeGenerator() method can be rewritten as follows
Hi! I’m glad you liked it! The snippet you provided should start from the number 2 because it is also a prime number.
Yeah, your absolutely right!
I guess my previous snippet would only be useful for people wanting to print all prime numbers except the even prime number (which is 2).
Here is my corrected code, which gives all the prime numbers:
Thank You,
Rishit
Thanks for this - it was super helpful.
I hacked it a bit and added a value limit for the prime number check, so that you can check the prime numbers up to X.
Changed up the prime generator function to have a limit (also set a default limit):
then used that limit in the call of the generator into a list:
so no longer need itertools and islice.
I’m glad that you find it useful 🙂
Nice post but 2 is not prime. You need to shift the n==2 if-statement to the top.
Hi Frank, I’m afraid 2 is a prime number
Hi Alejandro, yes, 2 is prime but your code says no.
The code returns True if it is a prime number. It says if n == 2, return True, which states that 2 is a prime number, because it is actually a prime number
Can u help me with describing the part of the code below i = 3?
Thanks.
Hi Jovan, the first ifs in the code cover the cases where i <= 1 and i == 2, so starting the iteration with i = 1 will repeat the first 2 cases. It's just a minimum performance improvement.