2007 JOI 春 Fermat (python)

素数についての謎の知識を二つ得た
1. p が素数のとき、x^n % p = x^(n%(p-1)) % p
2. p-1 > c > 1 のとき、x^n+y^n≡c (mod p) を満たす個数は０または x^n+y^n≡1 (mod p) を満たす個数に等しい。
証明はここをみてください。
フェルマー方程式/Fermat(JOI2007春合宿Day2 2nd) - Tsuzu's Notes

競プロにどの程度必要なのかは知らないが、2. を考えないと python では TLE し続けた。

p = int(input())
n = int(input())
n %= p-1
chk = [0]*p
for i in range(p):
  c = pow(i, n, p)
  chk[c] += 1
  
a0 = chk[0]**3
a1 = chk[0]*chk[1]*2*chk[1]
cnt = 0
for i in range(1,p):
  h = chk[i]
  if h == 0:
    continue
  a0 += h*chk[p-i]*chk[0]
  cnt += h//h
  
for i in range(2,p):
  a1 += chk[i]*chk[p-i+1]*chk[1]
  
print(a0+a1*cnt)