백준 17134번: 르모앙의 추측 (C++)

https://www.acmicpc.net/problem/17134

17134번: 르모앙의 추측

---

Level: Platinum I

Solved By: FFT

 

FFT 관련 문제 중 <다항식의 지수적 특성>을 활용한 문제입니다. (https://draw-code-boy.tistory.com/459) Golf Bot 문제의 아이디어와 비슷하기 때문에 이에 관해서는 링크를 참고하시기 바랍니다.

 

케이스마다 소수를 구하기에는 번거롭기 때문에 N의 최대 범위인 1,000,000까지는 소수를 모두 구해놓습니다. 그리고, 세미 소수라는 건 모든 소수에 2를 곱해주기만 하면 끝입니다. 즉, 소수와 세미 소수 간의 합을 모두 구하기 위해 다항식의 지수적 특성을 활용하여 소수에 해당하는 지수의 계수를 1로 두고, 세미 소수도 이러한 방식으로 다항식을 만들어 FFT를 통해 구한 식의 계수가 N에 대한 답이 됩니다.

#include <iostream>
#include <complex>
#include <algorithm>
#include <vector>
using namespace std;

int T;

const double PI = acos(-1);
typedef complex<double> cpx;

void FFT(vector<cpx> &f, bool inv = false){
    int n = f.size(), j = 0;
    vector<cpx> roots(n / 2);
    for(int i = 1; i < n; i++){
        int bit = (n >> 1);
        while(j >= bit){
            j -= bit;
            bit >>= 1;
        }
        j += bit;
        if(i < j) swap(f[i], f[j]);
    }
    double ang = 2 * PI / n * (inv ? -1 : 1);
    for(int i = 0; i < n / 2; i++){
        roots[i] = cpx(cos(ang * i), sin(ang * i));
    }
    for(int i = 2; i <= n; i <<= 1){
        int step = n / i;
        for(int j = 0; j < n; j += i){
            for(int k = 0; k < i / 2; k++){
                cpx u = f[j + k], v = f[j + k + i / 2] * roots[step * k];
                f[j + k] = u + v;
                f[j + k + i / 2] = u - v;
            }
        }
    }
    if (inv) for (int i = 0; i < n; i++) f[i] /= n;
}

vector<cpx> mul(vector<cpx> a, vector<cpx> b){
    FFT(a, 0);
    FFT(b, 0);
    
    // Inner Product
    for(int i = 0; i < a.size(); i++){
        a[i] = a[i] * b[i];
    }
    
    FFT(a, 1);
    
    vector<cpx> ret(a.size());
    for(int i = 0; i < a.size(); i++){
        ret[i] = cpx(round(a[i].real()), round(a[i].imag()));
    }
    return ret;
}

vector<int> getPrime(int n){
    vector<int> ret;
    vector<bool> visited(n + 1, false);
    for(int i = 2; i <= sqrt(n); i++){
        if(!visited[i]){
            ret.push_back(i);
            for(int j = i + i; j <= n; j += i){
                visited[j] = true;
            }
        }
    }
    
    for(int i = 2; i <= n; i++){
        if(!visited[i]) ret.push_back(i);
    }
    
    return ret;
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    auto p = getPrime(1000000);
    vector<cpx> f(1<<21), g(1<<21);
    
    for(int i = 0; i < p.size(); i++){
        f[p[i]] = cpx(1, 0);
        if(p[i] * 2 <= 1000000){
            g[2 * p[i]] = cpx(1, 0);
        }
    }
    
    auto res = mul(f, g);
    
    cin >> T;
    for(int t = 1; t <= T; t++){
        int N; cin >> N;
        
        cout << res[N].real() << '\n';
    }
    return 0;
}