Calculating a^b mod p for a large prime p

Question

I'm trying to write a python code that calculates a^b mod p, where p = 10^9+7 for a list of pairs (a,b). The challenge is that the code has to finish the calculation and output the result in < 1 second. I've implemented successive squaring to calculate a^b mod p quickly. Please see my code below:

from sys import stdin, stdout
rl = stdin.readline
wo = stdout.write
m = 10**9+7
def fn(a,n):
    t = 1
    while n > 0:
        if n%2 != 0: #exponent is odd
            t = t*a %m
        a = a*a %m
        n = int(n/2)
    return t%m

t = int(rl()) # number of pairs
I = stdin.read().split() # reading all pairs
I = list(map(int,I)) # making pairs a list of integers
# calculating a^b mod p. I used map because I read its faster than a for loop
s = list(map(fn,I[0:2*t:2],I[1:2*t:2])) 
stdout.write('\n'.join(map(str,s))) # printing output

for 200000 pairs (a,b) with a,b<10^9, my code takes > 1 second. I'm new to python and was hoping someone could help me identify the time bottle neck in my code. Is it reading input and printing output or the calculation itself? Thanks for the help!

Hi, Sorry this is my first time using stack overflow and didn't notice that someone responded. I tried using "pow" but still ran into the time issue. I decided to write my own code to solve the problem afterwards because I'm trying to learn python. It takes about 2 seconds, so 1 second more than the requirement. Here is a link to the problem: https://cses.fi/problemset/task/1095. Thank you @superbrain — Ibrahim M, Aug 07 '20 at 03:32

score 0 · Accepted Answer · answered Aug 07 '20 at 09:58

I don't see something wrong with your code from an efficiency standpoint, it's just unnecessarily complicated.

Here's what I'd call the straight-forward solution:

n = int(input())
for _ in range(n):
    a, b = map(int, input().split())
    print(pow(a, b, 10**9 + 7))

That did get accepted with PyPy3 but not with CPython3. And with PyPy3 it still took 0.93 seconds.

I'd say their time limit is inappropriate for Python. But try yours with PyPy3 if you haven't yet.

In case someone's wondering whether the map wastes time, the following got accepted in 0.92 seconds:

n = int(input())
for _ in range(n):
    a, b = input().split()
    print(pow(int(a), int(b), 10**9 + 7))

Thanks for your response @superbrain . I tried my code with PyPy3 and it got accepted with 0.35 seconds! I don't know the difference between CPython3 and PyPy3 so I'll do some reading today about both. Also, thank you fro sharing your code. Its very helpful to me to see cleaner and more direct solutions. Take care! — Ibrahim M, Aug 07 '20 at 14:32

Calculating a^b mod p for a large prime p

1 Answers1