백준 10531번: Golf Bot (C++)

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

10531번: Golf Bot

---

Solved By: FFT

 

문제가 영어로 되어있어서 간단히 요약하자면 N개의 정수가 주어지는데 뒤이어 주어지는 M개의 정수를 N개의 정수 중 1개 or 2개(중복 가능)를 골라 합하여 M개의 정수 중 몇 개를 나타낼 수 있는지에 대해 묻는 문제였습니다.

 

N개의 정수의 합을 Naive 하게 모두 구해놓고, M개의 정수를 확인하기에는 N개의 최대 개수가 200,000이라 O(N^2)로 TLE가 납니다.

 

FFT를 응용하는 두번째 문제이기도 했는데 이때까지 이 문제를 포함하여 FFT 응용에 대해서는 이런 식으로 응용할 수 있었습니다.

1. Convolution
2. 다항식 제곱의 지수적 특성(현 문제)

2번에 대하여 얘기하자면 다항식 X^3 + x^1을 제곱해봅시다. 이를 제곱하면 다항식은 X^(3 + 3) + 2X^(3 + 1) + X^(1 + 1)로 주어집니다. 즉, 다항식의 계수가 0이 아닌 수이며 지수가 각각 {3, 1}이 주어졌을 때, 제곱을 하면 {6, 4, 2}로 수를 중복 허용했을 때, 수 2개를 골라 나올 수 있는 합의 모든 경우를 구할 수 있는 것입니다.

 

즉, N개의 정수가 주어지면 N개의 정수를 이용하여 다항식을 만드는 것입니다. 다항식을 만드는 과정은 N개의 정수의 계수에 해당하는 차수에 계수를 1로 두면 됩니다. 만들어지는 다항식을 제곱하여 계수가 존재하는 항의 차수가 답이 됩니다.

 

하나로도 가능한지 물었기 때문에 M개의 정수가 N개의 정수에도 있는 수인지와 다항식 제곱의 지수에 해당하는 수의 계수 존재 여부를 따진다면 풀 수 있는 문제입니다.  

#include <iostream>
#include <complex>
#include <algorithm>
#include <vector>
using namespace std;

const double PI = acos(-1);
typedef complex<double> cpx;
bool one_only[200001];

void FFT(vector<cpx> &f, bool inv = false){
    int n = f.size(), j = 0;
    vector<cpx> 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(f[i], f[j]);
    }
    double ang = 2 * PI / n * (inv ? -1 : 1);
    for(int i = 0; i < n / 2; i++){
        roots[i] = cpx(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++){
                cpx u = f[j + k], v = f[j + k + i / 2] * roots[step * k];
                f[j + k] = u + v;
                f[j + k + i / 2] = u - v;
            }
        }
    }
    if (inv) for (int i = 0; i < n; i++) f[i] /= n;
}

vector<cpx> mul(vector<cpx> a, vector<cpx> b){
    int n = 1;
    while((n < a.size() + 1) || (n < b.size() + 1)) n *= 2;
    n *= 2;
    
    a.resize(n);
    b.resize(n);
    vector<cpx> c(n);
    
    FFT(a, 0);
    FFT(b, 0);
    
    // Inner Product
    for(int i = 0; i < n; i++){
        c[i] = a[i] * b[i];
    }
    
    FFT(c, 1);
    
    vector<cpx> ret(n);
    for(int i = 0; i < n; i++){
        ret[i] = cpx(round(c[i].real()), round(c[i].imag()));
    }
    return ret;
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    vector<cpx> polynomial(200001, cpx(0, 0));
    int n; cin >> n;
    for(int i = 0; i < n; i++){
        int v; cin >> v;
        one_only[v] = true;
        polynomial[v] = cpx(1, 0);
    }
    
    auto res = mul(polynomial, polynomial);
    
    int ans = 0;
    
    int m; cin >> m;
    for(int i = 0; i < m; i++){
        int v; cin >> v;
        if(res[v].real() != 0 || one_only[v] == true){
            ans++;
        }
    }
    cout << ans;
    return 0;
    
}