Infinite list of prime numbers using Python generators

alebian

Alejandro Bezdjian

Posted on September 14, 2019

Infinite list of prime numbers using Python generators

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


Enter fullscreen mode Exit fullscreen mode

Now, using our is_prime function we can do:



def prime_generator():
    n = 1
    while True:
        n += 1
        if is_prime(n):
            yield n


Enter fullscreen mode Exit fullscreen mode

And that's it! Just call the function and get elements from it:



generator = prime_generator()

for i in range(10):
    print(next(generator))


Enter fullscreen mode Exit fullscreen mode

Or create a list of the first N prime numbers:



from itertools import islice

array = [x for x in islice(prime_generator(), 10)]


Enter fullscreen mode Exit fullscreen mode

As you could see, the iterator definition is one of the shortest and simplest among all languages.

Buy Me A Coffee

💖 💪 🙅 🚩
alebian
Alejandro Bezdjian

Posted on September 14, 2019

Join Our Newsletter. No Spam, Only the good stuff.

Sign up to receive the latest update from our blog.

Related