백준 17134번 르모앙의 추측

문제 요약

홀수 N이 주어졌을 때, N을 홀수 소수 하나와 짝수 세미소수의 합으로 나타내는 경우의 수를 구하시오. 여기서 세미소수란 두 소수를 곱한 수를 말한다.

N은 최대 100만이며 테스트 케이스의 수는 최대 10만이다.

풀이

N이 주어지면 N보다 작은 모든 홀수 소수 p에 대해서 N-p가 짝수 세미소수인 p의 개수를 찾으면 됩니다.

세미소수는 임의의 소수에 2를 곱한 수입니다.

에라토스테네스의 체를 이용해서 짝수 세미소수와 홀수 소수들을 전처리하면 어떤 수 x가 주어졌을 때 x가 홀수 소수인지 그리고 N-x가 짝수 세미소수인지 판단하는 것은 쉽습니다. 그러나 이를 이용하면 테스트 케이스가 많아서 시간초과를 받습니다.

이를 해결하기 위해서 생성함수를 이용합니다.

다항식 $f(x)$와 $g(x)$를 만들 것인데 $f(x)$는 홀수 소수 차수의 항들의 계수가 1이고 나머지 항들의 계수는 0인 다항식입니다. $g(x)$는 짝수 세미소수 차수의 항들의 계수가 1이고 나머지는 0인 다항식입니다.

이제 $N$이 들어왔을 때 우리가 원하는 답은 $f(x), g(x)$를 곱한 다항식 $h(x)=f(x)g(x)$의 N차항의 계수가 됩니다.

$f$와 $g$의 차수는 100만으로 나이브하게 다항식의 곱을 구하면 시간초과가 납니다. FFT를 이용하면 이를 빠르게 구할 수 있고 문제를 해결할 수 있습니다.

#include<bits/stdc++.h>
using namespace std;
#define MAXN 1000000
typedef complex<double> cdbl;

bool chk[MAXN+5];
vector<int> primes;
void sieve() {
    for(int i=2;i<=MAXN;++i) {
        if(!chk[i]) {
            for(int j=i+i;j<=MAXN;j+=i) chk[j] = true;
        }
    }
    for(int i=2;i<=MAXN;++i) {
        if(!chk[i]) primes.push_back(i);
    }
}

void fft(vector<cdbl> &a, bool inv) {
    int n = a.size();
    // bit reversal
    for(int i=1,j=0;i<n;++i) {
        int bit = n>>1;
        while(!((j^=bit) & bit)) bit >>=1;
        if(i<j) swap(a[i],a[j]);
    }
    for(int i=1;i<n;i<<=1) {
        double x = inv? M_PI / i : -M_PI / i;
        cdbl w = cdbl(cos(x),sin(x));
        for(int j=0;j<n;j += i<<1) {
            cdbl p = cdbl(1,0);
            for(int k=0;k<i;++k) {
                cdbl tmp = a[i+j+k] * p;
                a[i+j+k] = a[j+k] - tmp;
                a[j+k] += tmp;
                p *= w;
            }
        }
    }
    if(inv) {
        for(int i=0;i<n;++i) a[i] /= n;
    }
}

// h = fg
void multiply(vector<cdbl> &f, vector<cdbl> &g, vector<cdbl> &h) {
    fft(f,false); fft(g,false);
    for(int i=0;i<f.size();++i) h[i] = f[i] * g[i];
    fft(h,true);
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    int T; cin >> T;
    sieve();
    vector<cdbl> f(1<<21), g(1<<21), ans(1<<21);
    for(auto p:primes) {
        f[p] = cdbl(1,0);
        if(p*2<=MAXN) g[2*p] = cdbl(1,0);
    }
    multiply(f,g,ans);
    while(T--) {
        int n;
        cin >> n;
        cout << round(ans[n].real()) << '\n';
    }
    return 0;
}