【SPOJ】Triple Sums【FFT】

传送门：【SPOJ】Triple Sums

题目分析：

首先我们不考虑$i<j<k$这个条件，构造多项式：
$\qquad\qquad\qquad Y = \sum x^{a_i}$
那么$ai+aj+ak=S$的个数即$x^{a_i+a_j+a_k=S}$的个数，等价于$Y^3中x^S$的系数。
然后我们考虑容斥：
$\qquad\qquad(\sum x)^3 = {\sum x^3+3\sum x^2y+6\sum xyz}$
$\qquad\qquad\sum x^2y = (\sum x^2)(\sum x)-\sum x^3$
$\qquad\qquad\sum xyz = \frac{(\sum x)^3 -3(\sum x^2)(\sum x)+2\sum x^3}{6}$

然后$FFT$处理下，这道题就做完啦。

my code：

#include <stdio.h>
#include <string.h>
#include <math.h>
#include <algorithm>
using namespace std ;

typedef long long LL ;

#define clr( a , x ) memset ( a , x , sizeof a )
#define cpy( a , x ) memcpy ( a , x , sizeof a )

const int MAXN = 200005 ;
const double pi = acos ( -1.0 ) ;

struct Complex {
    double r , i ;
    Complex () {}
    Complex ( double r , double i ) : r ( r ) , i ( i ) {}
    Complex operator + ( const Complex& t ) const {
        return Complex ( r + t.r , i + t.i ) ;
    }
    Complex operator - ( const Complex& t ) const {
        return Complex ( r - t.r , i - t.i ) ;
    }
    Complex operator * ( const Complex& t ) const {
        return Complex ( r * t.r - i * t.i , r * t.i + i * t.r ) ;
    }
} ;

void FFT ( Complex y[] , int n , int rev ) {
    for ( int i = 1 , j , k , t ; i < n ; ++ i ) {
        for ( j = 0 , k = n >> 1 , t = i ; k ; k >>= 1 , t >>= 1 ) j = j << 1 | t & 1 ;
        if ( i < j ) swap ( y[i] , y[j] ) ;
    }
    for ( int s = 2 , ds = 1 ; s <= n ; ds = s , s <<= 1 ) {
        Complex wn ( cos ( rev * 2 * pi / s ) , sin ( rev * 2 * pi / s ) ) , w ( 1 , 0 ) , t ;
        for ( int k = 0 ; k < ds ; ++ k , w = w * wn ) {
            for ( int i = k ; i < n ; i += s ) {
                y[i + ds] = y[i] - ( t = w * y[i + ds] ) ;
                y[i] = y[i] + t ;
            }
        }
    }
    if ( rev == -1 ) for ( int i = 0 ; i < n ; ++ i ) y[i].r /= n ;
}

int n ;
int num[MAXN] ;
Complex x1[MAXN] ;
Complex x2[MAXN] ;
Complex x3[MAXN] ;

void solve () {
    int x , n1 = 1 ;
    clr ( num , 0 ) ;
    for ( int i = 1 ; i <= n ; ++ i ) {
        scanf ( "%d" , &x ) ;
        num[x + 20000] ++ ;
    }
    while ( n1 <= 80000 ) n1 <<= 1 ;
    for ( int i = 0 ; i < n1 ; ++ i ) x1[i] = Complex ( num[i] , 0 ) ;
    for ( int i = 0 ; i < n1 ; ++ i ) {
        if ( i % 2 == 0 ) x2[i] = Complex ( num[i / 2] , 0 ) ;
        else x2[i] = Complex ( 0 , 0 ) ;
    }
    FFT ( x1 , n1 , 1 ) ;
    FFT ( x2 , n1 , 1 ) ;
    for ( int i = 0 ; i < n1 ; ++ i ) x3[i] = x1[i] * x1[i] * x1[i] ;
    for ( int i = 0 ; i < n1 ; ++ i ) x2[i] = x2[i] * x1[i] ;
    FFT ( x3 , n1 , -1 ) ;
    FFT ( x2 , n1 , -1 ) ;
    for ( int i = 0 ; i < n1 ; ++ i ) {
        LL ans = ( LL ) ( x3[i].r + 0.5 ) - 3 * ( LL ) ( x2[i].r + 0.5 ) ;
        if ( i % 3 == 0 ) ans += 2LL * num[i / 3] ;
        ans /= 6 ;
        if ( ans ) printf ( "%d : %lld\n" , i - 60000 , ans ) ;
    }
}

int main () {
    while ( ~scanf ( "%d" , &n ) ) solve () ;
    return 0 ;
}