IOICJ 102 No Title ( FFT )

Problem Description:
Given N integers A[ 0, N ), you want to know how many ordered tuples ( i, j, k ) satisfy:
1 ≤ i, j, k ≤ N and A[ i ] + A[ j ] = A[ k ]

Constraints:
1≤T≤10
3≤n≤1e5
−50000≤a1,a2,…,an≤50000

Sample Input:
3
3
1 2 3
4
2 2 2 2
4
-10 10 0 10

Sample Output:
3
0
15

Solution:
Let cnt[ i ] = count of appearance of value i in array A:
We could think of it as a polynomial function:
Sigma{ cnt[ i ] * x ^ i }, and we know that
cnt2 = Sigma{ cnt[ i ] * x ^ i } * Sigma{ cnt[ i ] * x ^ i }, would imply:
cnt2[ i ] = Sigma{ cnt[ j ] * cnt[ i - j ] * x ^ i }, which is exactly what we want.
However, we must map the original range [ -50000, 50000 ] to [ 0, 100000 ], then apply FFT. Note that precision handling is required.

Note:
NTT seems to be too slow to pass ...?

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

typedef long long ll;
typedef vector< int > vi;

int T;

const int MAXN = 262144 * 2;

typedef long double ld;
typedef complex< ld > cplx;
const ld PI = acos( -1.0 );
const cplx I( 0, 1 );
cplx omega[ MAXN + 1 ];
void pre_fft(){
  for( int i = 0; i <= MAXN; ++i )
    omega[ i ] = exp( i * 2 * PI / MAXN * I );
}

void fft( int n, cplx a[], bool inv = false ){
  int basic = MAXN / n;
  int theta = basic;
  for( int m = n; m >= 2; m >>= 1 ){
    int mh = m >> 1;
    for( int i = 0; i < mh; ++i ){
      cplx w = omega[ inv ? MAXN - ( i * theta % MAXN ) : i * theta % MAXN ];
      for( int j = i; j < n; j += m ){
        int k = j + mh;
        cplx x = a[ j ] - a[ k ];
        a[ j ] += a[ k ];
        a[ k ] = w * x;
      }
    }
    theta = ( theta * 2 ) % MAXN;
  }
  int i = 0;
  for( int j = 1; j < n - 1; ++j ){
    for( int k = n >> 1; k > ( i ^= k ); k >>= 1 );
    if( j < i ) swap( a[ i ], a[ j ] );
  }
  if( inv )
    for( i = 0; i < n; ++i )
      a[ i ] /= n;
}

int N;
int A[ MAXN + 1 ];

cplx cnt[ MAXN + 1 ];
cplx cnt1[ MAXN + 1 ];
cplx cnt2[ MAXN + 1 ];
cplx cnt12[ MAXN + 1 ];

void init(){
  cin >> N;
  for( int i = 0; i < N; ++i ){
    cin >> A[ i ];
    A[ i ] += 50000; // [ -50000, 50000 ] -> [ 0, 100000 ]
  }
  for( int i = 0; i <= MAXN; ++i ){
    cnt[ i ] = cplx( 0, 0 );
    cnt1[ i ] = cplx( 0, 0 );
    cnt2[ i ] = cplx( 0, 0 );
    cnt12[ i ] = cplx( 0, 0 );
  }
  for( int i = 0; i < N; ++i ){
    cnt[ A[ i ] ] += cplx( 1, 0 );
    cnt1[ A[ i ] ] += cplx( 1, 0 );
    cnt2[ A[ i ] ] += cplx( 1, 0 );
  }
  fft( MAXN / 2, cnt1 ); // convert to points
  fft( MAXN / 2, cnt2 );
  for( int i = 0; i < MAXN / 2; ++i ){ // point-wise product
    cnt12[ i ] = cnt1[ i ] * cnt2[ i ];
  }
  fft( MAXN / 2, cnt12, 1 ); // reverse point-wise product to coefficient
}

void solve(){
  ll ans = 0;
  for( int i = 0; i < MAXN / 2; ++i ){
    if( i - 50000 < 0 ) continue;
    ans += real( cnt12[ i ] ) * real( cnt[ i - 50000 ] ) + 1e-9;
  }
  cout << ans << endl;
}

signed main(){
  pre_fft();
  ios::sync_with_stdio( 0 );
  cin >> T;
  for( int ti = 0; ti < T; ++ti ){
    init();
    solve();
  }
  return 0;
}