SPOJTSUM三重和(FFT+容斥原理)

来源

SPOJ TSUM

问题描述

给定一个由N个不同整数组成的序列s。考虑从序列中选取三个不同的索引处的整数的所有可能和。对于每个可达到的和，输出产生该和的不同三元组的数量。

约束条件：
N ≤ 40000, |si| ≤ 20000

输入格式

输入的第一行包含一个整数N。接下来的N行每行包含序列s的一个元素。

输出格式

对于每个可能的和，按照以下格式输出：
sum_value : number_of_triples
较小的和值应先输出。

样例输入

5
-1
2
3
0
5

样例输出

1 : 1
2 : 1
4 : 2
5 : 1
6 : 1
7 : 2
8 : 1
10 : 1

本题的核心在于利用快速傅里叶变换（FFT）来高效计算多项式乘积，从而解决三数求和的问题。首先，根据给定的序列构建三个相同的多项式，其中指数代表数字，系数代表该数字出现的次数。通过计算这三个多项式的乘积，可以得到有序且可重复选择数字的所有结果。使用FFT可以显著提高计算效率。

然而，这种方法得到的是有序的三元组计数，而题目要求的是无序的三元组。因此，需要通过容斥原理来调整结果。具体来说，先减去所有两个数字相同的情况，再加回所有三个数字相同的情况，最后将结果除以3的阶乘（即6），以消除重复计数。

为了优化性能，可以采取一些措施减少不必要的计算，例如在多项式相乘时直接对点值进行操作，避免多次执行离散傅里叶变换（DFT）过程。

#include 
#include 
#include 
#include 
using namespace std;
#define MAXN 131072
const double PI = acos(-1.0);

struct Complex {
    double real, imag;
    Complex(double _real = 0, double _imag = 0) : real(_real), imag(_imag) {}
    Complex operator+(const Complex &cp) const {
        return Complex(real + cp.real, imag + cp.imag);
    }
    Complex operator-(const Complex &cp) const {
        return Complex(real - cp.real, imag - cp.imag);
    }
    Complex operator*(const Complex &cp) const {
        return Complex(real * cp.real - imag * cp.imag, real * cp.imag + cp.real * imag);
    }
    void setValue(double _real = 0, double _imag = 0) {
        real = _real; imag = _imag;
    }
};

int len;
Complex wn[MAXN + 1], wn_anti[MAXN + 1];

void FFT(Complex y[], int op) {
    for (int i = 1, j = len >> 1, k; i         if (i         k = len >> 1;
        while (j >= k) {
            j -= k;
            k >>= 1;
        }
        if (j     }
    for (int h = 2; h <= len; h <<= 1) {
        Complex Wn = (op == 1 ? wn[h] : wn_anti[h]);
        for (int i = 0; i             Complex W(1, 0);
            for (int j = i; j > 1); ++j) {
                Complex u = y[j], t = W * y[j + (h >> 1)];
                y[j] = u + t;
                y[j + (h >> 1)] = u - t;
                W = W * Wn;
            }
        }
    }
    if (op == -1) {
        for (int i = 0; i     }
}

Complex A[MAXN], B[MAXN];
double ans[MAXN];
int s[40100], cnt[80100];

int main() {
    for (int i = 0; i <= MAXN; ++i) {
        wn[i].setValue(cos(2.0 * PI / i), sin(2.0 * PI / i));
        wn_anti[i].setValue(wn[i].real, -wn[i].imag);
    }
    
    int n;
    scanf("%d", &n);
    for (int i = 0; i         scanf("%d", &s[i]);
        s[i] += 20000;
        ++cnt[s[i]];
    }
    
    len = MAXN;
    
    for (int i = 0; i <= 40000; ++i) {
        B[i].setValue(cnt[i]);
    }
    FFT(B, 1);
    for (int i = 0; i         A[i] = B[i] * B[i] * B[i];
    }
    FFT(A, -1);
    for (int i = 0; i         ans[i] = A[i].real;
    }
    
    memset(cnt, 0, sizeof(cnt));
    for (int i = 0; i         ++cnt[s[i] + s[i]];
    }
    for (int i = 0; i <= 80000; ++i) {
        A[i].setValue(cnt[i]);
    }
    for (int i = 80001; i         A[i].setValue(0);
    }
    FFT(A, 1);
    for (int i = 0; i         A[i] = A[i] * B[i];
    }
    FFT(A, -1);
    for (int i = 0; i         ans[i] -= 3 * A[i].real;
    }
    
    for (int i = 0; i         ++ans[s[i] + s[i] + s[i]];
    }
    
    for (int i = 0; i         long long tmp = (long long)(ans[i] / 6.0 + 0.5);
        if (tmp) {
            printf("%d : %lld\n", i - 60000, tmp);
        }
    }
    return 0;
}