Prime numbers which can be written as sum of the squares of two numbers x and y

Question

The problem is:

Given a range of numbers (x,y) , Find all the prime numbers(Count only) which are sum of the squares of two numbers, with the restriction that 0<=x<y<=2*(10^8)

According to Fermat's theorem :

Fermat's theorem on sums of two squares asserts that an odd prime number p can be   
expressed as p = x^2 + y^2 with integer x and y if and only if p is congruent to 
1 (mod4).

I have done something like this:

import math
def is_prime(n):
    if n % 2 == 0 and n > 2:
        return False
    return all(n % i for i in range(3, int(math.sqrt(n)) + 1, 2))

a,b=map(int,raw_input().split())
count=0
for i in range(a,b+1):
    if(is_prime(i) and (i-1)%4==0):
        count+=1
print(count)

But this increases the time complexity and memory limit in some cases.

Here is my submission result:

Can anyone help me reduce the Time Complexity and Memory limit with better algorithm?

Problem Link(Not an ongoing contest FYI)

Answer 1

Do not check whether each number is prime. Precompute all the prime numbers in the range, using Sieve of Eratosthenes. This will greatly reduce the complexity.

Since you have maximum of 200M numbers and 256Mb memory limit and need at least 4 bytes per number, you need a little hack. Do not initialize the sieve with all numbers up to y, but only with numbers that are not divisible by 2, 3 and 5. That will reduce the initial size of the sieve enough to fit into the memory limit.

UPD As correctly pointed out by Will Ness in comments, sieve contains only flags, not numbers, thus it requires not more than 1 byte per element and you don't even need this precomputing hack.

Answer 2

You can reduce your memory usage by changing for i in range(a,b+1): to for i in xrange(a,b+1):, so that you are not generating an entire list in memory.

You can do the same thing inside the statement below, but you are right that it does not help with time.

return all(n % i for i in xrange(3, int(math.sqrt(n)) + 1, 2))

One time optimization that might not cost as much in terms of memory as the other answer is to use Fermat's Little Theorem. It may help you reject many candidates early. More specifically, you could pick maybe 3 or 4 random values to test and if one of them rejects, then you can reject. Otherwise you can do the test you are currently doing.

Answer 3

First of all, although it will not change the order of your time-complexity, you can still narrow down the list of numbers that you are checking by a factor of 6, since you only need to check numbers that are either equal to 1 mod 12 or equal to 5 mod 12 (such as [1,5], [13,17], [25,29], [37,41], etc).

Since you only need to count the primes which are sum of squares of two numbers, the order doesn't matter. Therefore, you can change range(a,b+1) to range(1,b+1,12)+range(5,b+1,12).

Obviously, you can then remove the if n % 2 == 0 and n > 2 condition in function is_prime, and in addition, change the if is_prime(i) and (i-1)%4 == 0 condition to if is_prime(i).

And finally, you can check the primality of each number by dividing it only with numbers that are adjacent to multiples of 6 (such as [5,7], [11,13], [17,19], [23,25], etc).

So you can change this:

range(3,int(math.sqrt(n))+1,2)

To this:

range(5,math.sqrt(n))+1,6)+range(7,math.sqrt(n))+1,6)

And you might as well calculate math.sqrt(n))+1 beforehand.

To summarize all this, here is how you can improve the overall performance of your program:

import math

def is_prime(n):
    max = int(math.sqrt(n))+1
    return all(n % i for i in range(5,max,6)+range(7,max,6))

count = 0
b = int(raw_input())
for i in range(1,b+1,12)+range(5,b+1,12):
    if is_prime(i):
        count += 1
print count

Please note that 1 is typically not regarded as prime, so you might want to print count-1 instead. On the other hand, 2 is not equal to 1 mod 4, yet it is the sum of two squares, so you may leave it as is...