Ronak Jain
Ronak Jain

Reputation: 21

How to optimize and find output for large inputs?

For an input number N, I am trying to find the count of numbers of special pairs (x,y) such that the following conditions hold:

finally print the output of the count modulo 1000000007


Sample Input: 2

Sample Output: 2

Explanation:

  • (0,2) Since F(0)+F(2)=2 which is prime
  • (1,2) Since F(1)+F(2)=3 which is prime
  • (2,1) is not considered as (1,2) is same as (2,1)

My code is:

def mod(x,y,p):
    res=1
    x=x%p
    while(y>0):
        if((y&1)==1):
            res=(res*x)%p
        y=y>>1
        x=(x*x)%p
    return res

def sod(x):
    a=str(x)
    res=0
    for i in range(len(a)):
        res+=int(a[i])
    return res

def prime(x):
    p=0
    if(x==1 or x==0):
        return False
    if(x==2):
        return True
    else:
        for i in range(2,(x//2)+1):
            if(x%i==0):
                p+=1
        if(p==0):
            return (True)
        else:
            return(False)

n=int(input())
res=[]
for i in range (n+1):
    for j in range(i,n+1):
        if(prime(sod(int(i))+sod(int(j)))):
            if([i,j]!=[j,i]):
                if([j,i] not in res):
                    res.append([i,j])
count=len(res)
a=mod(count,1,(10**9)+7)
print(res)
print(a)

I expect the output of 9997260736 to be 671653298 but the error is code execution timed out.

Upvotes: 2

Views: 194

Answers (1)

h4z3
h4z3

Reputation: 5478

Already posted a bit too long comments, so changing it to answer:

When thinking of such problems, don't translate the problem directly to code but see what can you do only once or in different order.

As of now, you're doing N*N passes, each time calculating a sum of digits for x and y (not that bad) AND factoring each sum to check whether it's prime (really bad). That means for sum s you're checking whether it's prime s+1 times! (for 0+s, 1+(s-1), ..., (s-1)+1, s+0).

What you can do differently?

Let's see what we know:

  • Sum of digits is the same for many numbers.

  • Sum of sod(x) and sod(y) is the same for many values.

  • Number is prime during its 1st and nth check (and checking whether it's prime is costly).

So the best would be to calculate prime numbers only once, and each of the sum only once. But how to do that when we have many numbers?

Change the direction of thinking: get the prime number, split it into two numbers (sodx and sody), then generate x and y from those numbers.

Example:

Prime p = 7. That give possible sums as 0+7, 1+6, 2+5, 3+4.

Then for each sum you can generate a number, e.g. for N=101 and sod=1, you have 1, 10, 100, and for sod=2 you have 2, 11, 20, 101. You can possibly store this, but generating this should not be that bad.

Other optimisation:

You have to think how to limit generating prime numbers using your N:

given N with lenN digits (remember, lenN is ~log(N)), the biggest sum of digits possible is 9*lenN (for N consisting of only 9's). That means our sodx and sody are <= 9*lenN, so prime p = sodx + sody <= 18*lenN

Look: that means 18*lenN checking for whether number is prime vs N*N checks your previous algorithm had!

Upvotes: 1

Related Questions