Python prime number checker function

Question

I was bored, and I was watching a sci-fi show where people were starting out during a long systems check asking each other whether numbers were prime or not, sort of passing the time as they do their inspections.

The show is irrelevant, but I wondered if I could build a prime number checker in Python 3 that would be very quick. I know this has been done before, but I was curious whether I could do it on my own.

Now, this works fine for smaller numbers, and is pretty darn quick, a second or two. But for large numbers, it takes a really long time.

I'm not a coder by profession, though, so I'm pretty sure this could be improved. Any speed improvements are welcome.

Required imports, and the 'test' function which takes input of a number and spits out the results is below,

Currently, the results spit out a tuple, the first value either being Boolean or the number '0' (because the number '1' is a unique case that is not prime but you can't just say 'False' for), and a short string description.

from typing import AnyStr, Optional, Union

def is_prime(num: int) -> (Optional[Union[bool, int]], Optional[AnyStr]):
    if num <= 1:
        return 0, "Neither prime nor composite."

    if num <= 3:
        return True, "Prime"

    for i in range(2, num):
        if (num % i) == 0:
            return False, "Not Prime"

    return True, "Prime"

if __name__ == "__main__":
    n = int(input("Value: "))
    print(is_prime(n))

I'm curious how I can improve the performance of this code for larger numbers so that it's a fairly fast function, rather than taking a very long time to iterate over values. I know that the iterations in the for loop are taking the hugest amount of time, but I'm not a mathematician so I'm not sure the 'best' approach for speeding up this. Don't yell at me too harshly, this was more a quick 'mind test' to see if I could come up with something that worked.

Answer 1

^{Disclaimer: I can read Python, but I'm not a Python developer.}

First of all, your code follows PEP8, which is a plus (although a second empty line before is_prime is necessary). And it's great that you try static type hints.

However, there are two things amiss here. First of all, Optional isn't honest at that point. We always return a value. And we always return a str, not AnyStr. We should therefore change the type hints accordingly:

from typing import Union


def is_prime(num: int) -> (Union[bool, int], str):

If there was a path where we could accidentally return None, None or None, b"12343", a static checker should warn us. So no Optional or AnyStr.

Since the types are now out of a way, let's have a look at documentation before we head to algorithms. A docstring is a nice addition:

def is_prime(num: int) -> (Union[bool, int], str):
    """Checks whether the given number is prime."""

    if num <= 1:
    ...

You can probably come up with a more detailed one so that the code itself contains the meaning of True, False and 0.

Now to the algorithm. Your algorithm has \$\mathcal O(n)\$ worst time complexity. However, for any pair \$x\$,\$y\$, such that \$x < y\$ and \$ x y = n \$ we must have \$x \le \sqrt{n}\$. Therefore, we can stop at math.ceil(sqrt(n)), not n. This yields \$\mathcal O(\sqrt{n})\$ worst time complexity.

For very large numbers this will still be too slow. There are other algorithms for that, which are out of scope for a small review.

Answer 2

I'm not a mathematician, so I can't comment on 'correct' ways to identify primes - however one possible optimisation is to consider factors.

If a number isn't divisible by n, then it also won't be divisible by any other number for which n is a factor.

For example if a number isn't divisible by 2, then it also won't be divisible by any other even number, since all even numbers can be generated by 2 * something.

I have some ideas about how you could implement this in code, but I'll leave them out the answer for now, as this sounds like a learning example, and it might be an interesting challenge for you to consider them.

(I'm also hopeful that this observation will be surpassed by someone with a mathematical background giving some more info on how to definitively test for primes in an efficient manner!).

Answer 3

Well, it is correct to say that is_prime(1) = False. Do you really want to add all this complexity to your function just to handle the special case of n<2? It depends on what you're using it for. If you're just wanting to be able to give Rodney the correct answer, he's probably not giving n<2, and if he does, you can just say "not prime" without running the program. You should consider just raising an error when n<2, rather than complicating the output.

Answer 4

It's not a good idea, to return redundant information of a simple function. If you want to distinguish three cases, None, False and True would be good. If you also want a description, use a dictionary:

DESCRIPTION_OF_IS_PRIME_RESULT = {
    None: "Neither prime nor composite.",
    False: "Prime",
    True: "Not Prime",
}

Optimizations are, test only every second number and end at sqrt(num):

def is_prime(num):
    """ tests if num is a prime
    returns: None for num <= 1 otherwise
       True or False whether num is prime or not  
    """
    if num <= 1:
        return None

    if num % 2 == 0:
        return num == 2

    for i in range(3, int(sqrt(num)) + 1, 2):
        if num % i == 0:
            return False

    return True

if __name__ == "__main__":
    n = int(input("Value: "))
    print(DESCRIPTION_OF_IS_PRIME_RESULT[is_prime(n)])

Answer 5

Since other answers have covered style and efficiency of your code, I thought I'd contribute with another algorithm. The Miller-Rabin primality test is a probabilistic primality test which is fast. Really fast. I have implemented it myself in both Java and C++, and it could handle 100-digit numbers in a few milliseconds.

I found a Python 2.7 version of it here, and with some very minor changes it works with python 3.x:

from random import randrange

def miller_rabin(n, k=10):
    if n == 1:
        return False
    if n == 2 or n == 3:
        return True
    if not n & 1:
        return False

    def check(a, s, d, n):
        x = pow(a, d, n)
        if x == 1:
            return True
        for i in range(s - 1):
            if x == n - 1:
                return True
            x = pow(x, 2, n)
        return x == n - 1

    s = 0
    d = n - 1

    while d % 2 == 0:
        d >>= 1
        s += 1

    for i in range(k):
        a = randrange(2, n - 1)
        if not check(a, s, d, n):
            return False
    return True

I have not written this code, but it seems to be working after some slight modifications (before the modifications it would hang on n = 1, and give an error for n = 3).