백준 17134번 - 르모앙의 추측

www.acmicpc.net/problem/17134

17134번: 르모앙의 추측

 

#include <iostream>
#include <complex>
#include <vector>
#include <cstring>
#include <string>
#include <algorithm>
 
// Koosaga 팀노트 FFT 코드
using namespace std;
typedef complex<double> base;
typedef long long ll;
const double PI = acos(-1);
 
void fft(vector<base>& a, bool inv = false) {
    int n = a.size(), j = 0;
    // 분할 정복 기법, 이때 n은 무조건 2^n 승임
    vector<base> 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(a[i], a[j]);
    }
    double ang = 2 * acos(-1) / n * (inv ? -1 : 1);
    for (int i = 0; i < n / 2; i++) {
        roots[i] = base(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++) {
                base u = a[j + k], v = a[j + k + i / 2] * roots[step * k];
                a[j + k] = u + v;
                a[j + k + i / 2] = u - v;
            }
        }
    }
    if (inv) for (int i = 0; i < n; i++) a[i] /= n;
}
 
vector<ll> multiply(vector<ll>& v, vector<ll>& w) {
    vector<base> fv(v.begin(), v.end()), fw(w.begin(), w.end());
    // n이 무조건 2^n 이여야 하기 때문에 변환!
    int n = 2; while (n < v.size() + w.size()) n <<= 1;
    fv.resize(n); fw.resize(n);
    fft(fv, 0); fft(fw, 0);
    for (int i = 0; i < n; i++) fv[i] *= fw[i];
    fft(fv, 1);
    vector<ll> ret(n);
    for (int i = 0; i < n; i++) ret[i] = (ll)round(fv[i].real());
    return ret;
}
 
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
 
    // 1. 에라토스테네스의 체로 100만 까지의 모든 소수 구하기 
    bool prime[1000004];
    memset(prime, true, sizeof(prime));
    for (int i = 2; i <= 1000000; i++) {
        if (prime[i] == true) {
            for (int k = 2; i * k <= 1000000; k++) {
                prime[i * k] = false;
            }
        }
    }
 
    // 2. 하나의 배열에는 홀수인 소수만 (a), 다른 배열에는 모든 소수*2 저장(b)
    // -> 이 때 그냥 저장하는게 아님! ex) 소수 = 3 이라면 a[3] = 1 이런식으로
    // 이유 : 다항식의 곱셉을 통해 차수를 사용하려고 
 
    vector<ll> a(1000004), b(1000004);
    for (int i = 2; i <= 1000000; i++) {
        if (prime[i] == true) { // i가 소수 라면
            // 일단 b에는 모든 소수*2(세미소수)를 저장
            if(i*2 <= 1000000) b[2 * i] = 1;
            // i가 홀수라면 a에 저장
            if (i % 2 == 1) a[i] = 1;
        }
    }
 
    // 3. FFT 사용
    vector<ll> ret = multiply(a, b);
 
    // 4. ret[n] 값이 정답!
    int t, n;
    cin >> t;
    while (t--) {
        cin >> n;
        cout << ret[n] << "\n";
    }
}