백준 BOJ 20176번 : Needle 풀이 (FFT, 고속 푸리에 변환)

백준 BOJ 20176번 : Needle

문제 난이도 : Platinum I

알고리즘 분류 : FFT (고속 푸리에 변환)

 

 

3개의 직선이 있고 3개의 직선 위에 점들의 좌표가 주어질 때 서로 다른 직선 위의 3개의 점이 일직선에 있는 쌍이 몇 개나 있는지를 구하는 문제입니다.

입력은 3개의 직선에 대해 직선 위의 점의 수 N, 점의 x 좌표 N개가 순서대로 주어집니다.

 

2
1 2
3
0 1 2
2
2 3

 

예를 들어 입력이 위와 같이 주어졌다고 가정하면, 맨 위에서 좌표 1, 중간에서 좌표 2, 맨 아래에서 좌표 3을 선택하면 하나의 직선을 만들 수 있습니다.

또한 맨 위에서 좌표 2, 중간에서 좌표 2, 맨 아래에서 좌표 2를 선택해도 역시 직선 하나를 만들 수 있습니다.

따라서 위의 입력에 대한 답은 2가지, 즉 2가 출력되어야 합니다.

 

이제 문제에 대한 이해가 되었으면 풀어봅시다.

 

 

 

먼저 세 점이 일직선에 있으려면, (세 직선의 간격은 같으므로) 세 직선 위의 점의 좌표의 간격이 동일해야 합니다.

즉, 세 직선 위의 점의 좌표를 위에서부터 순서대로 x, y, z라고 하면 y = (x+z)/2가 성립되어야 하며, 이는 양변에 2를 곱해 x + z = 2y로 만들어줄 수 있습니다.

 

FFT를 활용하면, 여러 개의 x와 z들에 대하여 x + z = k를 만족하는 경우의 수를 찾을 수 있습니다.

(좌표를 배열 내에 저장할 수 있을 정도의, 즉 몇 백만 정도의 작은 범위 내에서여야 합니다. 그 이상의 범위에 대해서는 잘 모르겠습니다.)

 

 

 

예를 들어 x에서 구멍이 뚫린 좌표들이 0부터 나열했을 때 2, 6, ...이고 z에서 구멍이 뚫린 좌표들이 0, 3, 4, 6, ...이라고 한다면 위와 같이 x = (0, 0, 1, 0, 0, 0, 1, 0, ... , 1), z = (1, 0, 0, 1, 1, 0, 1, 0, ... , 0)와 같이 나타낼 수 있을 것입니다.

 

두 벡터에 FFT를 수행하여 이산 곱을 구해주면, 각 좌표에 대해 나타나는 숫자는, 두 수를 더해서 해당 좌표 값을 만들 수 있는 경우의 수가 됩니다.

예를 들어 xz 벡터에서 좌표 6에 해당하는 값은 2인데, 이는 x의 2과 z의 4를 더해서 만들거나, x의 6과 z의 0을 더해서 만들 수 있기 때문입니다.

 

 

 

y의 좌표 k를 지나면서 만들 수 있는 직선의 수는, y_k * xz_2k와 같이 표현할 수 있습니다.

그리고 이러한 좌표는 문제에서 주어진 범위인 -30000 ~ +30000이므로 모두 검사해서 더해주어야 합니다.

FFT를 한 번 수행하면 모든 좌표에 대한 xz_2k 값이 구해져있기 때문에 이 과정에서의 시간 복잡도는 고려할 필요 없습니다. (어차피 O(60000)이므로)

 

문제가 되는 부분은 문제에서 주어지는 좌표는 -30000부터라는 것인데, 이는 좌표 전체를 +30000씩 옮겨주면 됩니다.

즉, 값을 받을 때 모두 30000씩 더해서 벡터에 저장해주고, 검사 범위도 0 ~ 60000으로 잡아주면 됩니다.

 

그럼 이제 코드로 구현해봅시다.

 

 

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

const double PI = acos(-1);
typedef complex<double> cpx;

void FFT(vector<cpx> &v, bool inv) {
    int S = v.size();

    for(int i=1, j=0; i<S; i++) {
        int bit = S/2;

        while(j >= bit) {
            j -= bit;
            bit /= 2;
        }
        j += bit;

        if(i < j) swap(v[i], v[j]);
    }

    for(int k=1; k<S; k*=2) {
        double angle = (inv ? PI/k : -PI/k);
        cpx w(cos(angle), sin(angle));

        for(int i=0; i<S; i+=k*2) {
            cpx z(1, 0);

            for(int j=0; j<k; j++) {
                cpx even = v[i+j];
                cpx odd = v[i+j+k];

                v[i+j] = even + z*odd;
                v[i+j+k] = even - z*odd;

                z *= w;
            }
        }
    }

    if(inv)
        for(int i=0; i<S; i++) v[i] /= S;
}

vector<cpx> mul(vector<cpx> v, vector<cpx> u) {
    int S = 2;
    while(S < v.size() + u.size()) S *= 2;

    v.resize(S); FFT(v, false);
    u.resize(S); FFT(u, false);

    vector<cpx> w(S);
    for(int i=0; i<S; i++) w[i] = v[i] * u[i];
    FFT(w, true);

    return w;
}

main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL), cout.tie(NULL);

    vector<cpx> v(60001), u(60001), w(60001);

    int N; cin >> N;
    while(N--) {
        int x; cin >> x;
        v[x + 30000] = cpx(1, 0);
    }

    int M; cin >> M;
    while(M--) {
        int x; cin >> x;
        u[x + 30000] = cpx(1, 0);
    }

    int K; cin >> K;
    while(K--) {
        int x; cin >> x;
        w[x + 30000] = cpx(1, 0);
    }

    vector<cpx> vw = mul(v, w);

    int ans = 0;
    for(int i=0; i<=60000; i++)
        ans += round(u[i].real()) * round(vw[i*2].real());

    cout << ans << "\n";
}

 

역시나 FFT 코드 부분은 건들 필요가 없고, main 함수 부분에서의 구현만 조금 적어주면 됩니다.

값을 입력받은 뒤 좌표를 30000만큼 더해서 벡터에 위치를 저장해주고, FFT를 수행하여 convolution을 구해줍니다.

이후 위에서 구한 식인 u[i].real() * vw[i*2].real()의 sum을 구해주기만 하면 됩니다.

 

만약 문제의 조건에 의해 같은 좌표가 여러 번 들어올 수 있다면 cpx의 real 값을 1 더해주는 방식으로 갱신해주면 됩니다.