10531번: Golf Bot - FFT

거듭제곱끼리의 곱셈은 차수의 덧셈과 같다. 여기서 차수는 \(k_i\)와 같으므로, FFT를 이용해 다항 방정식의 곱셈을 빠르게 계산하면 문제를 해결할 수 있다.

 

\(k_i\)차의 계수를 1로 하는 방정식을 2개 만든다.

예제 입력에서 \(k_i\)가  1, 3, 5이므로, \(1 * x^1 + 1 * x^3 + 1* x^5\)의 다항 방정식을 2개 만든다.

두 다항 방정식을 곱해서 나온 방정식의 어떠한 차수의 계수가 1일 때, 그 차수가 \(d_i\)라면 \(d_i\)는 답이 될 수 있다.

예를 들어 \(d_i\)가 4일 경우 \(x^1 * x^3 = x^4 (1 + 3 = 4)\)이므로 4는 답이 될 수 있다.

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

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef tuple<int, int, int> tiii;
typedef complex<double> cd;

const double PI = acos(-1);
const int SIZE = 200001;
int N, M;
int GB[SIZE];

void FFT(vector<cd>& p, cd w)
{
    int n = p.size();
    if(n == 1) return;
    
    vector<cd> even, odd;
    for(int i=0; i<n; i++)
    {
        if(i % 2) odd.push_back(p[i]);
        else even.push_back(p[i]);
    }
    
    FFT(even, w * w), FFT(odd, w * w);
    
    cd cw{1, 0};
    for(int i=0; i<n/2; i++)
    {
        p[i] = even[i] + cw * odd[i];
        p[i + n/2] = even[i] - cw * odd[i];
        cw *= w;
    }
}

vector<cd> mul(vector<cd>& F, vector<cd>& G)
{
    int n = 1;
    while(n < F.size() || n < G.size()) n *= 2;
    n *= 2;
    F.resize(n), G.resize(n);
    
    cd w = {cos(2 * PI / n), sin(2 * PI / n)};
    FFT(F, w), FFT(G, w);
    
    vector<cd> FG;
    for(int i=0; i<n; i++) FG.push_back(F[i] * G[i]);
    
    cd rw = {cos(-2 * PI / n), sin(-2 * PI / n)};
    FFT(FG, rw);
    for(int i=0; i<n; i++) FG[i] /= n;
    
    return FG;
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    
    cin >> N;
    vector<cd> F(SIZE, 0), G(SIZE, 0);
    
    for(int i=1; i<=N; i++)
    {
        int gb; cin >> gb;
        F[gb] = G[gb] = GB[gb] = 1;
    }
    
    vector<cd> FG = mul(F, G);
    
    cin >> M;
    int ans = 0;
    while(M--)
    {
        int d; cin >> d;
        if(GB[d] == 1 || int(FG[d].real()+0.5) != 0) ans++;
    }
    cout << ans;
}