[BOJ 5051] 피타고라스의 정리

$a^{2}+b^{2} \equiv c^{2} \pmod{n}$을 만족하는 $(a,\ b,\ c)$($a \le b$) 쌍의 개수를 구하는 것이 목표이다.

집합 $\{s\}$를 1부터 $n-1$까지의 자연수의 제곱을 $n$으로 나눈 나머지의 집합이라 정의하자.

집합 $s$의 원소를 각 항의 차수로 하고 항의 계수가 1인 다항식 $P(x)$를 생각하자.

e.g. $n=7$ 이면 $P(x)=2x+2x^{2}+2x^{4}$

$P(x)^{2}$를 했을 때 차수가 집합 $s$에 포함되어있는 것들의 계수를 전부 더하면 된다. 반면 문제의 조건에서 $a \le b$가 있다. 따라서 위 더한 값에서 $a=b$인 경우를 제외시킨 뒤 2로 나누고 $a=b$인 경우를 다시 더하면 된다. $a=b$인 경우는 $P(x)$의 차수를 2배 해준 뒤 $n$으로 나누어 집합 $s$에 존재하는 개수를 세면 된다. 

 

FFT를 제외한 과정은 모두 $O(n/2)$에 해결할 수 있다.

 

# 코드 보기

#define _USE_MATH_DEFINES
#include <math.h>
#include <complex>
#include <vector>
#include <algorithm>
using namespace std;
 
#define sz(v) ((int)(v).size())
#define all(v) (v).begin(),(v).end()
typedef long long ll;
typedef complex<double> base;
const int SZ = 1<<20;
 
void fft(vector <base> &a, bool invert)
{
    int n = sz(a);
    for (int i=1,j=0;i<n;i++){
        int bit = n >> 1;
        for (;j>=bit;bit>>=1) j -= bit;
        j += bit;
        if (i < j) swap(a[i],a[j]);
    }
    for (int len=2;len<=n;len<<=1){
        double ang = 2*M_PI/len*(invert?-1:1);
        base wlen(cos(ang),sin(ang));
        for (int i=0;i<n;i+=len){
            base w(1);
            for (int j=0;j<len/2;j++){
                base u = a[i+j], v = a[i+j+len/2]*w;
                a[i+j] = u+v;
                a[i+j+len/2] = u-v;
                w *= wlen;
            }
        }
    }
    if (invert){
        for (int i=0;i<n;i++) a[i] /= n;
    }
}
 
int n;
 
void multiply(const vector<ll> &a,vector<ll> &res)
{
    vector <base> fa(all(a));
    fft(fa,false);
    for (ll i=0;i<SZ;i++) fa[i] *= fa[i];
    fft(fa,true);
    for (ll i=0;i<SZ;i++) res[i%n] +=ll(fa[i].real()+(fa[i].real()>0?0.5:-0.5));
}
 
int main(){
    ll i,j,ans=0,mu=0;
    vector<ll> a(SZ),res(SZ),b(SZ),c(SZ);
    scanf("%lld",&n);
    for(i=1;i<=n/2;i++){
        if(i!=n-i){
            a[i*i%n]+=2;
            c[2*i*i%n]+=2;
        }
        else{
            a[i*i%n]++;
            c[2*i*i%n]++;
        }
        b[i]=i*i%n;
    }
    multiply(a,res);
    for(i=1;i<=n/2;i++){
        if(i!=n-i) ans+=res[b[i]]*2;
        else ans+=res[b[i]];
    }
    for(i=1;i<=n/2;i++){
        if(i!=n-i) mu+=c[b[i]]*2;
        else mu+=c[b[i]];
    }
    ans=(ans-mu)/2+mu;
    printf("%lld",ans);
}

# 닫기