백준 10531번 Golf Bot

먼저 수가 N개 들어옵니다. 들어오는 수의 범위는 1이상 200,000이하입니다.

그리고 수가 M개 들어옵니다. 문제에서 원하는 것은 N개의 수들 중에서 1개 혹은 2개를 사용해서 합해서 M개의 수 중에서 만들 수 있는 수들의 개수입니다.

N개의 수들로 다항식을 만듭니다. 등장한 수를 차수로 가지는 항의 계수만 1이고 나머지 항들의 계수는 0입니다.

이 다항식의 제곱을 구하면 수 2개의 합으로 구할 수 있는 수들을 알 수 있습니다. 이를 통해 문제 해결이 가능합니다. 다항식의 곱을 FFT로 구하지 않으면 시간제한에 걸립니다.

#include<iostream>
#include<complex>
#include<vector>
#include<cmath>
using namespace std;
typedef complex<double> cdbl;

bool possible[200005];
int N,M;

void fft(vector<cdbl> &a, bool inv) {
    int n = a.size();
    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;
}

void multiply(vector<cdbl> &a, vector<cdbl> &b, vector<cdbl> &res) {
    vector<cdbl> pa = a, pb = b;
    fft(pa,false); fft(pb,false);
    for(int i=0;i<a.size();++i) res[i] = pa[i]*pb[i];
    fft(res,true);
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    vector<cdbl> p(1<<19);
    cin >> N;
    int tmp;
    for(int i=0;i<N;++i) {
        cin >> tmp;
        possible[tmp] = true;
        p[tmp] = cdbl(1,0);
    }
    vector<cdbl> res(1<<19);
    multiply(p,p,res);
    int ans = 0;
    cin >> M;
    for(int i=0;i<M;++i) {
        cin >> tmp;
        if(possible[tmp] || round(res[tmp].real())) ++ans;
    }
    cout << ans << '\n';
}