Take the 2-minute tour ×
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. It's 100% free, no registration required.

I have the following function that checks if a number is prime:

def is_prime(n):
    """ Is n prime? """

    if n < 2:
        return False
    if n in (2, 3):
        return True
    if n % 2 == 0 or n % 3 == 0:
        return False

    limit   = int(math.sqrt(n))
    divisor = 5

    while divisor <= limit:
        if n % divisor == 0 or n % (divisor + 2) == 0:
            return False
        divisor += 6

    return True

Most of it are checks and calculations to save up time, to make it faster than a simple loop through a range.

Can it be written in a more pythonic way? Can it be improved in some way?

Any other feedback is welcome too.

share|improve this question
    
If you want to save time you may want to use a faster, probabilistic test, like Baillie-PSW or Miller-Rabin for larger inputs (usually you quickly check for divisibility by some small primes, then move on to them, but it depends how your function will be used). –  Thomas Jul 1 at 3:18
    
@Thomas I have no idea of how to implement something like that. I know probabilistic tests exist, but I don't know enough to make use of them, I don't even know where to start. –  JCPedroza Jul 1 at 11:13
add comment

4 Answers 4

up vote 11 down vote accepted

Your code looks nice, seems to follow PEP 8 and has docstrings. I don't have much to say about it but here are a few changes to make it look more pythonic.

The while loop can be rewritten using a for loop with range (or xrange).

for divisor in xrange(5, 1+int(math.sqrt(n)), 6):
    if n % divisor == 0 or n % (divisor + 2) == 0:
        return False
return True

Then, this looks like a scenario where we could use all/any. The obvious way to convert the code above is :

return not any(n % divisor == 0 or n % (divisor + 2) == 0 for divisor in xrange(5, int(math.sqrt(n))+1, 6))

but it looks better (and makes more sense) to use all with reversed conditions :

return all(n % divisor != 0 and n % (divisor + 2) != 0 for divisor in xrange(5, int(math.sqrt(n))+1, 6))

then, you can use the fact that non-zero numbers evaluates to True in a boolean context (and zero evaluates to False) :

return all(n % divisor and n % (divisor + 2) for divisor in xrange(5, int(math.sqrt(n))+1, 6))

Then, you could use the same kind of idea in the previous if then merge duplicated code :

return n % 2 and n % 3 and all(n % divisor and n % (divisor + 2) for divisor in xrange(5, int(math.sqrt(n))+1, 6))

This is pretty much all I have to say about the code. Now, I think it might be worth commenting the algorithm a bit more.

def is_prime(n):
    """ Check if n (assumed to be an int) is a prime number. """
    # Handle values up to 3
    if n < 2:
        return False
    if n in (2, 3):
        return True
    # Handles bigger values : divisors for n have to be
    # * multiples of 2
    # * multiples of 3
    # * of the form 5 + 6*i or 5 + 6*i + 2 (other cases already handled)  
    return n % 2 and n % 3 and all(n % divisor and n % (divisor + 2) for divisor in xrange(5, int(math.sqrt(n))+1, 6))
share|improve this answer
6  
Super redundant comments. is_prime """Check if it's a prime""" What else would a method called "is_prime" do? Also # Handle values up to 3 followed by if n < 2 etc, not necessary when the code is that straight forward. Comments are good up to a point, but should be describing concepts or operations that aren't apparent. You have just documented self-documenting code. –  user11177 Jun 30 at 17:20
1  
Whoah, "return n % 2 and n % 3 and all(n % divisor and n % (divisor + 2) for divisor in xrange(5, int(math.sqrt(n))+1, 6))" I always aim for readable code and nothing like this :X –  RvdK Jul 1 at 10:24
3  
This I think is a prime example of how you can obfuscate great, easy to read code in an attempt to write a one-liner (which also violates PEP8 by being 115 chars long) Not pythonic in my book but more.. perlithic. –  Voo Jul 1 at 11:13
add comment
if n < 2:
    return False
if n in (2, 3):
    return True

These can be combined:

if n < 4:
    return n in (2, 3)

This:

divisor = 5

while divisor <= limit:
    if n % divisor == 0 or n % (divisor + 2) == 0:
        return False
    divisor += 6
return True

Really wants to be either a for loop:

for divisor in xrange(5, limit+1, 6):
    if n % divisor == 0 or n % (divisor + 2) == 0:
        return False
return True

or a generator:

return all(n % divisor and n % (divisor + 2) for divisor in xrange(5, limit+1, 6))
share|improve this answer
add comment

This can definitely be improved by using any one of of a multitude of prime number algorithms, e.g the Sieve of Atkin or the Sieve of Eratosthenes.

However, without having you implement one of those algorithms, your code looks to be pretty well optimized.

I have a couple suggestions that will help your code look and feel more Pythonic:

  1. Take advantage of the fact that 0 == False. So instead of n % 2 == 0 we can do not n % 2. We can then logically factor out the not to make your if statements like this:

    if not(n % 2 and n % 3):
    ...
    if not(n % divisor and n % (divisor + 2)):
    
  2. Don't tab align equal signs. This is actually a Pythonic convention pet peeve:

    # No
    limit   = int(math.sqrt(n))
    divisor = 5
    
    # Yes
    limit = int(math.sqrt(n))
    divisor = 5
    
  3. Its good you use docstrings. They can be really helpful. However, I would suggest making yours more descriptive instead of asking a question.

  4. I would probably use num instead of n as its slightly less ambiguous.

share|improve this answer
1  
Your first point is great for Golfing, but seems to hurt readability. Maybe it's just a Python thing that I'm not used to seeing. –  ckuhn203 Jul 1 at 1:08
1  
If you do use not like that, you should leave a space before (, because it isn't a function call. –  otus Jul 2 at 9:28
add comment

I would change this:

while divisor <= limit:
    if n % divisor == 0 or n % (divisor+2) == 0:
        return False
    divisor += 6

To this:

for pivot in range(6,limit+1,6):
    if n % (pivot-1) == 0 or n % (pivot+1) == 0:
        return False
share|improve this answer
    
I see the appeal of initializing to 6 if you're going to be incrementing by 6. However, now divisor never acts as a divisor, so it probably should be renamed. –  200_success Jun 30 at 15:23
    
@200_success: I certainly agree, thanks :) –  barak manos Jun 30 at 15:28
    
I reckon limit+2 is what you should be using :-) (Also, I suggest you'd have a look at the other answers and they somehow apply to your code - this doesn't make your answer wrong in any way but it is just for your personal information). –  Josay Jun 30 at 16:59
add comment

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.