DEV Community šŸ‘©ā€šŸ’»šŸ‘Øā€šŸ’»

DEV Community šŸ‘©ā€šŸ’»šŸ‘Øā€šŸ’» is a community of 968,547 amazing developers

We're a place where coders share, stay up-to-date and grow their careers.

Create account Log in
Alejandro Bezdjian
Alejandro Bezdjian

Posted on • Updated on

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

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.

Buy Me A Coffee

Top comments (12)

Collapse
 
orenovadia profile image
orenovadia

Since the time complexity of is_prime is O(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 to N (this is because sieves use previously found primes to eliminate other composite numbers, instead of performing an exhaustive search for each number).

Collapse
 
rishitc profile image
Rishit Chaudhary

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

def primeGenerator():
    n = 3
    while True:
        if isPrime(n):
            yield n
        n += 1
Collapse
 
alebian profile image
Alejandro Bezdjian Author

Hi! Iā€™m glad you liked it! The snippet you provided should start from the number 2 because it is also a prime number.

Collapse
 
rishitc profile image
Rishit Chaudhary

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:

def primeGenerator():
    n = 2
    while True:
        if isPrime(n):
            yield n
        n += 1

Thank You,
Rishit

Collapse
 
billybobza profile image
Billy Einkamerer • Edited on

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.

Collapse
 
alebian profile image
Alejandro Bezdjian Author

Iā€™m glad that you find it useful šŸ™‚

Collapse
 
frankzielen profile image
Frank

Nice post but 2 is not prime. You need to shift the n==2 if-statement to the top.

Collapse
 
alebian profile image
Alejandro Bezdjian Author

Hi Frank, Iā€™m afraid 2 is a prime number

Collapse
 
frankzielen profile image
Frank

Hi Alejandro, yes, 2 is prime but your code says no.

Thread Thread
 
b_wil34 profile image
Bwillis

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

Collapse
 
jovan_cosovic profile image
Jovan Cosovic

Can u help me with describing the part of the code below i = 3?
Thanks.

Collapse
 
alebian profile image
Alejandro Bezdjian Author

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.

šŸŒš Friends don't let friends browse without dark mode.

Sorry, it's true.