FFT

문제 : https://www.acmicpc.net/problem/10531

10531번: Golf Bot

 

개요

FFT는 고속 푸리에 변환을 의미하며, PS에서는 보통 다항식의 계산을 빠르게 계산할 때에 사용합니다. 알고리즘의 핵심 아이디어는  (n+1)개의 점을 모두 지나는 유일한 n차 다항식을 구할 수 있다는 것으로, 직접 다항식을 곱하지 않고 (n+1)개의 점을 복소수의 특징을 활용하여 빠르게 구하고 이를 활용하여 다항식을 구하는 알고리즘입니다.

 

소스코드

#include <stdio.h>
#include <math.h>
#include <vector>
#include <complex>
#define NMAX 200010
using namespace std;
typedef complex< double > cpx;
 
int N, M, t, ret;
vector< int > q, check;
vector< cpx > dist(NMAX);
 
int sz;
const double PI = acos(-1);
 
void FFT(vector< cpx > &v, cpx w) {
    // 종료 조건
    int n = v.size();
    if(n == 1) return;
 
    // 홀수차항과 짝수차항 나누기
    vector< cpx > even, odd;
    for(int i=0;i<n;i++) {
        if(i%2 == 0) even.push_back(v[i]);
        else odd.push_back(v[i]);
    }
 
    // 홀수차항과 짝수차항으로 나눠서 구하기
    FFT(even,  w*w);
    FFT(odd, w*w);
 
    // 따로 구한 값 합치기
    cpx wp(1.0, 0);
    for(int i=0;i<n/2;i++) {
        v[i] = even[i] + wp*odd[i];
        v[i + n/2] = even[i] - wp*odd[i];
 
        wp *= w;
    }
}
 
vector< int > mul(vector< cpx > a, vector< cpx > b) {
    // 2^n <= sz 구하기 및 정규화
    for(sz=1;sz<NMAX;sz*=2);
    sz *= 2;
 
    a.resize(sz);
    b.resize(sz);
 
    // FFT >> (sz*2)개의 점 구하기
    cpx w(cos(2*PI/sz), sin(2*PI/sz));
    FFT(a, w);
    FFT(b, w);
 
    // 결과 계산하기
    vector< cpx > c;
    for(int i=0;i<sz;i++) {
        c.push_back(a[i]*b[i]);
    }
 
    // 켤레복소수를 활용한 iDFT
    cpx _w(w.real(), -w.imag());
    FFT(c, _w);
 
    // 결과 리턴
    vector< int > ret;
    for(int i=0;i<sz;i++) {
        c[i] /= cpx(sz, 0);
        ret.push_back(round(c[i].real()));
    }
 
    return ret;
 
}
 
int main() {
    // 초깃값
    dist[0] = cpx(1, 0);
 
    // 입력
    scanf("%d", &N);
    for(int i=1;i<=N;i++) {
        scanf("%d", &t);
        dist[t] = cpx(1, 0);
    }
    scanf("%d", &M);
    for(int i=1;i<=M;i++) {
        scanf("%d", &t);
        q.push_back(t);
    }
 
    // FFT
    check = mul(dist, dist);
 
    // check[q[i]] > 0 : q[i]거리로 이동할 수 있음
    for(int i=0;i<M;i++) ret += (check[q[i]]>0);
 
    // 결과 출력
    printf("%d", ret);
}