AtCoder Problems in C++ #fermat - フェルマー方程式 (Fermat)

2007年 日本情報オリンピック春合宿OJから、その提出コードの解答例です。
AtCoderとは、コンテストを通じて、プログラミングやアルゴリズムを学習するサービスです。
atcoder.jp
プログラミングコンテストとは、「与えられた問題をいかに素早く、正確に」解くことができるかを競うものです。
「競技プログラミング」を略して、「競プロ」などと呼ばれています。

#fermat - フェルマー方程式 (Fermat)

https://www.ioi-jp.org/camp/2007/2007-sp-tasks/2007-sp-day2_21.pdf

僕が作成、提出したコードは、以下のとおりです。

/*
AtCoder Problems in C++
#fermat - フェルマー方程式 (Fermat)
https://www.ioi-jp.org/camp/2007/2007-sp-tasks/2007-sp-day2_21.pdf
提出コードの解答例
https://neguse-atama.hatenablog.com
*/
#include<bits/stdc++.h>
using namespace std;
using ll=long long;
ll mod;
ll f(ll x,ll y){
  ll a=x;
  ll b=1;
  while(y>1){
    if(y%2==1){
      b=b*a%mod;
    }
    a=a*a%mod;
    y/=2;
  }
  return a*b%mod;
}
int main(void){
  ll p,n;
  cin>>p>>n;
  mod=p;
  ll a[p],c[p]={},ans=0;
  for(ll i=0;i<p;i++){
    a[i]=f(i,n);
    c[a[i]]++;
  }
  for(ll i=0;i<p;i++){
    ans+=c[(a[i]*2)%mod];
    for(ll j=i+1;j<p;j++){
      ans+=c[(a[i]+a[j])%mod]*2;
    }
  }
  cout<<ans<<endl;
  return 0;
}