The next prime number in Python

Question

I have written a solution to a DM::OJ exercise to find the next Prime Number after a given number (input). I'm seeking to optimize it so that it will run for any number up to 1000000000 in 2 seconds or less.

Here is my code:

from functools import reduce

num = int(input())

def factors(n):    
  return set(reduce(list.__add__, 
  ([i, n//i] for i in range(1, int(n**0.5) + 1) if n % i == 0)))

f = (factors(num))

if (len(f) != 2): 
  while (len(f) != 2): 
    num += 1
    f = factors(num)
  print (num)
else: 
  print (num)

Just as a side note, I also found another way to execute this exact same thing, but that one is slower, so I wouldn't mind bothering with that way. However, just for reference though, I'm still adding the slow one anyways:

num = int(input())

def factors(x):
  a = []
  for i in range(1, x + 1):
      if x % i == 0:
          a.append(i)
  return a

f = (factors(num))

if (len(f) != 2): 
  while (len(f) != 2): 
    num += 1
    f = factors(num)
  print (num)
else: 
  print (num)

Answer 1

if (len(f) != 2): 
  while (len(f) != 2): 
    num += 1
    f = factors(num)
  print (num)
else: 
  print (num)

In places where you use both if and while, they can usually be collapsed into a single while. You can think of an if as a single use while. In your first if block, the final print(num) would only get executed once len(f) becomes 2, so it can just be left outside of the while loop. Which will make something like this:

while (len(f) != 2): 
  num += 1
  f = factors(num)
print (num)

---

As for ways of speeding up, you could consider keeping a list of previous primes and only dividing by those. That will cut down on the comparisons in the factor search. It might require some prior calculation, but the next calculations will be faster.

Another thing that could help is taking steps bigger than +1 while looking for the next prime. We know that all even numbers bigger than 2 are not prime, for example, and also that all numbers divisible by 3 that are bigger than 3 are not prime. So out of every 6 numbers, only two are possibly prime. If we use a counter variable n, then 6n, 6n+2, 6n+3, and 6n+4 are definitely not prime, so you could just check 6n+1 and 6n+5. Although that would require rearranging some of your program.

Answer 2

There are a lot of primality tests out there. Calculating all factors and then checking the length of that list to be 2 is one of the slower ones.

A faster one is just to test if any number smaller than \$\sqrt{n}\$ divides the number evenly. If it does, the number is not prime and you can move on to the next number:

from math import sqrt

def is_prime(n):
    if n in (2, 3, 5, 7, 11):  # special case small primes
        return True
    if n % 2 == 0 or n == 1:  # special case even numbers and 1
        return False
    for i in range(3, int(sqrt(n)) + 1, 2):
        if n % i == 0:
            return False
    return True

Then you can get all numbers starting with num using itertools.count and skip all which are not prime using filter:

from itertools import count

if __name__ == "__main__":
    num = int(input())
    next_prime = next(filter(is_prime, count(num)))
    print(next_prime)

For the upper bound of \$2\cdot 10^9\$ given in the problem statement this takes about 2ms on my machine.

If this was not fast enough you could first have run a probabilistic test like Miller-Rabin to weed out some composite numbers and only run this for the once passing it.