백준 20176번 Needle

문제 요약

길이가 최대 5만인 수열 $A, B, C$가 주어집니다. 이 때 ,$A_i, B_j, C_k$가 등차수열이 되는 $(i,j,k)$ triplet의 개수를 구해야 합니다.

들어오는 수열의 원소들의 범위는 [-30,000, 30,000]입니다.

풀이

어떤 세 수 $a, b, c$가 등차수열을 이룬다면 아래와 같은 식이 성립합니다.
$$
2b=a+c
$$
따라서, 문제에서 원하는 것은 $2B_j=A_i+C_k$인 $(i,j,k)$의 개수입니다. 다행히도 세 수열 $A,B,C$의 각 원소의 범위는 크기가 6만입니다. 수열 $cnt_A(i)$를 $A$에서 원소가 $i$인 것의 개수라고 합시다.

그러면 $cnt_A$와 $cnt_C$의 Convolution $cnt$를 구하면 $cnt(k)$는 아래와 같은 의미를 가집니다.
$$
cnt(k)=(A_i+C_j=k\text{인 (i,j)의 개수})
$$

따라서, $cnt$와 $cnt_B$를 구한 뒤에 $cnt(2k) * cnt_B(k)$를 전부 구해주면 답이 됩니다.

#include<bits/stdc++.h>
using namespace std;
using ll = int;
using cdbl = complex<double>;
using ld = double;
const int START = 30000;
void fft(vector<cdbl> &a, bool inv) {
    int n = a.size();
    // bit reversal
    for(int i=1,j=0;i<n;++i) {
        int bit = n>>1;
        while(!((j^=bit) & bit)) bit >>=1;
        if(i<j) swap(a[i],a[j]);
    }
    for(int i=1;i<n;i<<=1) {
        double x = inv? M_PI / i : -M_PI / i;
        cdbl w = cdbl(cos(x),sin(x));
        for(int j=0;j<n;j+=i<<1) {
            cdbl p = cdbl(1,0);
            for(int k=0;k<i;++k) {
                cdbl tmp = a[i+j+k] * p;
                a[i+j+k] = a[j+k] - tmp;
                a[j+k] += tmp;
                p *= w;
            }
        }
    }
    if(inv) {
        for(int i=0;i<n;++i) a[i] /= n;
    }
}

vector<ll> multiply(vector<ll> &f, vector<ll> &g) {
    vector<cdbl> pf(f.begin(), f.end()), pg(g.begin(), g.end());
    int n = 1; while (n < f.size() + g.size()) n <<= 1;
    pf.resize(n); pg.resize(n);
    fft(pf, false); fft(pg, false);
    for (int i = 0; i < n; ++i) pf[i] *= pg[i];
    fft(pf, true);
    vector<ll> ret(n);
    for (int i = 0; i < n; ++i) {
        ret[i] = (ll)round(pf[i].real());
    }
    return ret;
}

int N, M, L;
vector<ll> A(60005), B(60005);
int v[50005];
int main() {
  cin.tie(nullptr); ios::sync_with_stdio(false);
  cin >> N;
  for(int i=0;i<N;++i) {
      int x; cin >> x; A[START+x]++;
  }
  cin >> M;
  for(int i=0;i<M;++i) {
      int x; cin >> v[i];
  }
  cin >> L;
  for(int i=0;i<L;++i) {
      int x; cin >> x; B[START+x]++;
  }
  auto ret = multiply(A,B);
  ll ans = 0;
  for(int i=0;i<M;++i) ans += ret[2*START+2*v[i]];
  cout << ans << '\n';
  return 0;
}