[JOI精進録] 難易度7「fermat」(2007年日本情報オリンピック春合宿3-2) - [blog

問題リンク

問題概要

$x^n+y^n \equiv z^n \pmod p$ を満たす $x,y,z$ $(0 \le x,y,z \le p-1)$ の組はいくつ考えられますか?

解説

AC解法1 (100/100点)

繰り返し二乗法を使って $O(p \; log \; n)$ で $n$ 乗の数を $p$ で割ったあまりを列挙し、
$A_i=$ $n$ 乗して $p$ で割ったあまりが $i$ になる $x$ の数
を数え上げます。
$x$ と $y$ について $A$ を全探索すると、 $z$ は定数時間で求められるため $O(P^2)$ で答えることができます。

ACコード1

#include <bits/stdc++.h>
using namespace std;

int main(){
  int p,n; cin >> p >> n;
  long long r(0);
  vector<int> A(p);
  for (int i(0);i < p;++i){
    int r = 1;
    for (int t = n,u = i;t;t>>=1,u*=u,u%=p) if ((t&1)) r*=u,r%=p;
    ++A[r];
  }
  for (int i(0);i < p;++i) for (int k(0);k < p;++k){
    r += A[i]*A[k]*A[(i+k)%p];
  }
  cout << r << endl;
}

提出結果1

atcoder.jp

AC解法2(1より高速) (100/100点)

$A$ を1では全探索したところでFFTです。ACLに定義されているのを使うと簡単に実装できます。

ACコード2

#include <bits/stdc++.h>
#include <atcoder/convolution>
using namespace std;

int main(){
#define int long long
  int p,n,r; cin >> p >> n;
  vector<int> A(p);
  for (int i(0);i < p;++i){
    r = 1;
    for (int t = n,u = i;t;t>>=1,u*=u,u%=p) if ((t&1)) r*=u,r%=p;
    ++A[r];
  }
  vector<int> B = atcoder::convolution_ll(A,A);
  for (int i(0);i < B.size();++i) B[i]*=A[i%p];
  cout << accumulate(B.begin(),B.end(),0ll) << endl;
}

提出結果2

atcoder.jp