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!

Answer 1

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))