[C++] 10531번 Golf Bot - 수학, FFT

10531번 Golf Bot

1. 문제

• 골프 기계가 할 수 있는 거리 모음 N개의 정수와, 떨어진 구멍까지의 거리 모음 M개의 정수가 주어진다. 두 번 쳐서 공이 들어가는 구멍의 개수는 몇 개인가? 단, 구멍을 향하여 정면으로만 칠 수 있으며 뒤로 치는 것은 불가능하다.
  • 입력
    • -1- : 정수 \(N(1 \leq N \leq 200,000)\)
    • -N- : 공을 보낼 수 있는 거리 \(k_j(1 \leq k_j \leq 200,000)\)
    • -1- : 정수 \(M(1 \leq M \leq 200,000)\)
    • -M- : 구멍까지의 거리 \(d_j(1 \leq d_j \leq 200,000)\)
  • 출력
    • -1- : 공이 들어갈 수 있는 구멍의 개수

 

2. 재정의

• N개의 정수중 2개를 중복 가능하게 뽑아서 합이 M개의 정수들 중 같을 수 있는 정수의 개수를 구하여라.

3. 해결 방법

• 비거리를 차수로 갖는 항으로 이루어진 다항식 f를 만들자. 여기에 상수항 1을 더한 후 제곱하면 두 비거리의 합을 차수로 가지는 다항식이 만들어지며, 상수항으로 인해 한 비거리의 차수를 갖는 항도 만들어진다. 따라서 두 다항식을 곱하였을 때 구멍까지의 거리를 차수로 가지는 항이 존재하면 구멍에 넣는 것이 가능함을 의미한다.
• 다항식의 곱은 FFT로 빠르게 수행하자.

4. 실수한 점, 개선할 점

• 상수항을 넣어줌으로써 1번 치는 경우와 2번 치는 경우를 나누지 않고 한 번에 계산할 수 있었다.
• FFT를 처음에는 재귀방법으로 구현하였다가 비재귀적으로 푸는 방법으로 계선하였다. 이 코드는 다른 분 것을 참고하였으며 아래 링크를 달았다. 자세한 FFT 구현방법은 아래 링크를 참고하자.

 

<코드>

#include <iostream>
#include <vector>
#include <complex>

using namespace std;

// 복소수 자료구조
typedef complex<double> cpx;
// 복소수 벡터 자료구조
typedef vector<cpx> cvec;
// 정수 벡터 자료구조
typedef vector<int> vi;

const double pi = acos(-1);

const int MAX_N = 200001;

// <문제>
// 비거리 개수, 구멍의 개수, 구멍까지의 거리, 비거리
int N, M, d, k;
// 들어갈 수 있는 구멍의 개수
int cnt;

// 다항식 계수
cvec a(MAX_N);

// Fast Fourier Transform
// 함수 f, 역연산 inv
void FFT(cvec &f, bool inv) {
    int n = (int)f.size();
    cvec ftmp(f);
    for(int i=0; i<n; i++) {
        // bs: binary scale
        int bs = n, shift = 0, idx = i;
        while(bs > 1) {
            if(idx & 1) shift += bs >> 1;
            idx >>= 1;
            bs >>= 1;
        }
        f[shift + idx] = ftmp[i];
    }
    
    for(int i=1; i<n; i <<= 1) {
        double x = inv ? -pi/i : pi/i;
        cpx w(cos(x), sin(x));
        for(int j=0; j<n; j += i << 1) {
            cpx wp(1, 0);
            for(int k=0; k<i; k++) {
                cpx tmp = f[i + j + k] * wp;
                f[i + j + k] = f[j + k] - tmp;
                f[j + k] += tmp;
                wp *= w;
            }
        }
    }
    
    if(inv)
        for(int i=0; i<n; i++)
            f[i] /= n;
}

// 두 다항식의 곱 반환
cvec pol_mul(cvec a, cvec b) {
    int n=1;
    
    //cout << a.size();
    while((n <= a.size()) || (n <= b.size())) n <<= 1;
    n <<= 1;
    a.resize(n); b.resize(n);
    
    FFT(a, 0); FFT(b, 0);
    
    for(int i=0; i<n; i++)
        a[i] *= b[i];
    
    FFT(a, 1);
    
    return a;
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    
    cin >> N;
    for(int i=0; i<N; i++) {
        cin >> k;
        a[k] = cpx(1, 0);
    } a[0] = 1;
    // 비거리차수로 구성된 다항식의 제곱을 구하자
    cvec b(a);

    // 다항식의 곱
    cvec c = pol_mul(a, b);
    
    cin >> M;
    for(int i=0; i<M; i++) {
        cin >> d;
        if(round(c[d].real()))
            cnt++;
    }
    
    cout << cnt;
    
    return 0;
}

 

<문제 바로가기>

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

10531번: Golf Bot

 

<FFT의 원리>

https://justicehui.github.io/hard-algorithm/2019/09/04/FFT/

FFT in PS

 

<C++ 복소수 자료구조, complex 클래스>

http://seismic.yonsei.ac.kr/complex.html

http://seismic.yonsei.ac.kr/complex.html

 

<비재귀적 FFT 구현>

https://tistory.joonhyung.xyz/6

Fast Fourier Transform

 

※이제 알고리즘을 본격적으로 공부하기 시작해서 아직 초보입니다. 혹시 제가 잘못 알고 있는 점이나 더 좋은 풀이 방법이 있어 댓글에 남겨주시면 감사히 하나하나 열심히 읽어보겠습니다. 좋아요, 단순한 댓글 한마디라도 저에겐 큰 힘이 됩니다! 감사합니다.