Wonseok Shin | BOJ 20176, ICPC Seoul 2020 Needle

Contents

풀이
코드

풀이

결국 문제의 요점은, 위 / 중간/ 아래를 $u_i, m_i, l_i$ 라고 할 때, $u_i + l_j = 2m_k$ 인 순서쌍 $(i, j, k)$의 개수를 세는 문제이다.

이 문제는 매우 잘 알려진 Convolution 문제로, 수열을 카운팅 다항식으로 인코딩하면 쉽게 풀 수 있다. 수열이 1, 3, 4, 5, 5면 $x + x^3 + x^4 + 2x^5$ 인 식으로. 이렇게 인코딩해서 $u(x), m(x), l(x)$ 를 만든 다음,

$u(x)l(x)$을 계산하면 그 $n$차항의 계수가 $u_i + l_j = n$ 인 순서쌍 $i, j$ 의 개수가 된다. 따라서, 각 $n$에 대해, $m_k = n$인 $k$의 개수와 $u(x)l(x)$의 $2n$차 항의 계수를 곱해서 더하면 그대로 구하는 답이 되고, 이를 FFT를 이용하여 $O(n \log n)$에 구할 수 있음이 잘 알려져 있다.

여담으로… ICPC 서울 리저널은 이런 Do you know *** 문제가 꽤 많이 나오는 편인 것 같다. 2017 FFT, 2017 인예 FFT, 2019 인예 LiChaoTree / LR Flow …

코드

#include <bits/stdc++.h>
#pragma GCC optimize("O3")
#pragma GCC optimize("Ofast")
#pragma GCC target("avx,avx2,fma")
#define ll long long
#define int ll
#define eps 1e-7
#define all(x) ((x).begin()),((x).end())
#define usecppio ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
using namespace std;
using pii = pair<int, int>;

#define sz(v) ((int)(v).size())
typedef vector<int> vi;
typedef complex<double> base;

void fft(vector <base> &a, bool invert)
{
    int n = sz(a);
    for (int i=1,j=0;i<n;i++){
        int bit = n >> 1;
        for (;j>=bit;bit>>=1) j -= bit;
        j += bit;
        if (i < j) swap(a[i],a[j]);
    }
    for (int len=2;len<=n;len<<=1){
        double ang = 2*M_PI/len*(invert?-1:1);
        base wlen(cos(ang),sin(ang));
        for (int i=0;i<n;i+=len){
            base w(1);
            for (int j=0;j<len/2;j++){
                base u = a[i+j], v = a[i+j+len/2]*w;
                a[i+j] = u+v;
                a[i+j+len/2] = u-v;
                w *= wlen;
            }
        }
    }
    if (invert){
        for (int i=0;i<n;i++) a[i] /= n;
    }
}

void multiply(const vi &a,const vi &b,vi &res)
{
    vector <base> fa(all(a)), fb(all(b));
    int n = 1;
    while (n < max(sz(a),sz(b))) n <<= 1;
    fa.resize(n); fb.resize(n);
    fft(fa,false); fft(fb,false);
    for (int i=0;i<n;i++) fa[i] *= fb[i];
    fft(fa,true);
    res.resize(n);
    for (int i=0;i<n;i++)
        res[i] = (int)(fa[i].real()+(fa[i].real()>0?0.5:-0.5));
}

vi u(121212, 0), m(121212, 0), l(121212, 0), c;

void input(vi &v)
{
    int n; cin >> n;
    for (int i = 0; i < n; i++)
    {
        int x; cin >> x;
        v[x+30000]++;
    }
}

int32_t main()
{
    usecppio
    input(u); input(m); input(l);
    multiply(u, l, c);
    int ans = 0;
    for (int i = 0; i <= 60000; i++)
        ans += m[i] * c[2*i];
    cout << ans << '\n';
}