Python中的模幂运算

Question

I was trying to solve the ZSUM problem at SPOJ using Python 2.7.9 as my coding language and designed a program to do so. As the code runs perfectly but gives a TLE at the judge, I guess it is not fast enough. Is there is possible to optimize the below code to meet the judge requirement or it is impossible to beat the challenge using Python.

Link to the problem: http://www.spoj.com/problems/ZSUM/

def zsum(n,k):
  a=2*pow(n-1,k,10000007)
  b=pow(n,k,10000007)
  c=2*pow(n-1,n-1,10000007)
  d=pow(n,n,10000007)
  zsum=(a+b+c+d)%10000007
  print zsum

def main():
  while True:
    n,k=map(int,raw_input().split())
    if n==k==0:
        break
    else:
        zsum(n,k)
 main()

Answer 1

I don't know if this will help (if you read the problem's comments you will see someone saying that it is impossible to solve in Python - that can happen on online judges with slower languages) but you can optimize the code:

def zsum(n,k):
  a=2*pow(n-1,k,10000007)   # (1)
  b=pow(n,k,10000007)       # (2)
  c=2*pow(n-1,n-1,10000007) # (1)
  d=pow(n,n,10000007)       # (2)
  zsum=(a+b+c+d)%10000007
  print zsum

Note that in (1) , you are computing pow(n - 1, min(k, n - 1)) twice. You could compute it once and then use modular exponentiation only for what is left. Same for (2) .

Answer 2

Since there are zero accepted python solutions in a total of 2335 successful submissions, i think, no matter how much you optimise your solution, it is unlikely to manage to get accepted with python. Although python is a very useful language, it is not preferred in programming contests, as it is extremely slow (compared to C/C++ for instance). If you know how to code in C++ you should definitely give it a shot, although you must write you own modular exponentiation procedure.