ICPC Seoul Regional 2020 풀어보기 - 2

총 12문제 중, 난이도 순서대로 H, I, L, D, F, K를 풀어보자.

---

H - Needle

벽이 1 간격으로 세 레이어가 있는데, 각 벽에는 구멍이 나있다고 한다. 각 벽의 구멍 위치가 주어졌을 때, 바늘이 지나갈 수 있는 경우의 수를 구하라는 문제다.

 

구멍의 개수는 한 줄당 최대 5만개, 위치값은 -30000부터 30000의 정수이다. 전체탐색하면 위쪽 벽 50000개, 아래쪽 벽 50000개를 조합해 평균값을 내고 그 값을 가지는 중간 벽 구멍이 있으면 풀 수 있지만 50000*50000 = 2.5x10^9이므로 시간 초과가 난다. n^2이 아닌 nlogn으로 풀어야 한다.

 

키포인트

위쪽 벽과 아래쪽 벽 위치 값을 평균낸다는 아이디어를 재활용하자. 총 60001개의 위치가 가능하므로 구멍의 위치는 60001차원 벡터로 표현 가능하다. 그리고 중간 벽의 n번째 구멍을 지날 수 있기 위해선 위쪽 벽에 (n-k), 아래쪽 벽에 n+k 위치에 구멍이 있어야 한다. 즉 더해서 2n이 되어야한다.

이를 각 계수가 구멍의 존재 유무(1 또는 0인) 60000차 다항식 곱셈으로 쉽게 표현할 수 있다. 중간 벽 5 위치에 구멍이 뚫려있다고 하자. 이를 확인하기 위해서 위아래 벽 다항식을 곱하면 5차*5차 + 4차*6차 + 3차*7차가 되고 구멍 유무가 저절로 곱해진다. 이를 nlogn으로 푸는 방법인 FFT를 사용할 수 있다. 보통 쿨리-튜키 알고리즘을 쓴다.

 

주의할 점이 한가지 있는데, 위치가 음수일 수 있으므로 배열에 입력받을 때 65536x2짜리 배열의 인덱스 30000을 0으로 하자. FFT를 진행하면 두 방정식이 곱해지므로 60000이 0의 위치가 된다. iDFT 결과 위치에 30000을 빼면 원래대로 30000을 0점으로 사용할 수 있다.

 

코드

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

typedef complex<double> cxd;

int off = 30000;
const double PI = 3.14159265358979323846;
void FFT(vector<cxd> &a, bool inv){

    //FFT Code by https://aruz.tistory.com/1

    int n = a.size();
    if(n==1) return;

    vector<cxd> ev(n/2), od(n/2);
    for(int i=0 ; 2*i<n ; i++){
        ev[i] = a[2*i];
        od[i] = a[2*i+1];
    }
    FFT(ev,inv);
    FFT(od,inv);

    double euler = 2 * PI/n * (inv?-1:1);
    cxd w(1), wn(cos(euler),sin(euler));

    for(int i=0 ; 2*i<n ; i++, w*=wn){
        a[i] = ev[i]+w*od[i];
        a[n/2+i] = ev[i]-w*od[i];
    }

    if(inv) for(cxd &e : a) e/=2;    // divide by N = 2^k 
}

vector<cxd> up(65536*2);
vector<int> md(65536);
vector<cxd> dn(65536*2);
int U,M,D;

int main(){
    ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
    cin >> U;
    for(int i=0;i<U;i++){
        int t;
        cin >> t;
        up[t+off]={1,0};
    }
    cin >> M;
    for(int i=0;i<M;i++){
        int t;
        cin >> t;
        md[t+off]=1;
    }
    cin >> D;
    for(int i=0;i<D;i++){
        int t;
        cin >> t;
        dn[t+off]={1,0};
    }
    
    
    vector<int> total;
    FFT(up,0);
    
    FFT(dn,0);
    for(int i=0;i<65536*2;i++){
        up[i]*=dn[i];
    }
    FFT(up,1);
    int count = 0;
    for(int i=0;i<65536*2;i++){
        int real = round(up[i].real());
        if(i%2==0){
            if(md[i/2]==1) count+=real;
        }
    }
    cout << count;
    //code
}