SPOJ-TSUM Triple Sums

Description

You're given a sequence s of N distinct integers.
Consider all the possible sums of three integers from the sequence at three different indicies.
For each obtainable sum output the number of different triples of indicies that generate it.
**
Constraints:**
N <= 40000, |si| <= 20000

Input

The first line of input contains a single integer N.
Each of the next N lines contain an element of s.

Output

Print the solution for each possible sum in the following format:
sum_value : number_of_triples

Smaller sum values should be printed first.

Example

Input:
5-
1
2
3
0
5
Output:
1 : 1
2 : 1
4 : 2
5 : 1
6 : 1
7 : 2
8 : 1
10 : 1

Explanation:
4 can be obtained using triples ( 0, 1, 2 ) and ( 0, 3, 4 ).
7 can be obtained using triples ( 0, 2, 4 ) and ( 1, 3, 4 ).

题意

求所有3个数相加可能的和，及该和所对应的\((i,j,k)\)个数,\(i \le j \le k\)

题解

如果不考虑重复选择,我们直接对所有数加上20000,直接fft即可.

考虑有重复,我们要减去一个数选择两遍和一个数选择三遍的情况,一个数选择两遍同样可以通过fft求得,一个数选择三遍可以在输入时直接统计得出.

如果选择两个一样的数,那么第三个数有3种放置方式,所以要减去3倍的\(A * B\),同时,对于三个数都相同的情况,也在减去3倍\(A*B\)的时候被减掉了3次,实际上这个只需要减一次,我们再加回来两次即可

由于题目要求有序,答案还需要除以\(3!\)

代码

#include <bits/stdc++.h>
using namespace std;
const double pi = acos(-1.0);
typedef long long ll;
struct cp {
    double r, i;
    cp(double r = 0, double i = 0): r(r), i(i) {}
    cp operator + (const cp &b) {
        return cp(r + b.r, i + b.i);
    }
    cp operator - (const cp &b) {
        return cp(r - b.r, i - b.i);
    }
    cp operator * (const cp &b) {
        return cp(r * b.r - i * b.i, r * b.i + i * b.r);
    }
};
void change(cp a[], int len) {
    for (int i = 1, j = len / 2; i < len - 1; i++) {
        if (i < j) swap(a[i], a[j]);
        int k = len / 2;
        while (j >= k) {
            j -= k;
            k /= 2;
        }
        if (j < k) j += k;
    }
}
void fft(cp a[], int len, int op) {
    change(a, len);
    for (int h = 2; h <= len; h <<= 1) {
        cp wn(cos(-op * 2 * pi / h), sin(-op * 2 * pi / h));
        for (int j = 0; j < len; j += h) {
            cp w(1, 0);
            for (int k = j; k < j + h / 2; k++) {
                cp u = a[k];
                cp t = w * a[k + h / 2];
                a[k] = u + t;
                a[k + h / 2] = u - t;
                w = w * wn;
            }
        }
    }
    if (op == -1) {
        for (int i = 0; i < len; i++) {
            a[i].r /= len;
        }
    }
}
 
 
const int N = 2e5 + 50;
const int M = 2e4;
 
ll ans[N];
cp a[N], b[N];
ll num[N];
 
int main() {
    int n;
    scanf("%d", &n);
    int len1 = 1;
    for (int i = 0; i < n; i++) {
        int x;
        scanf("%d", &x);
        len1 = max(len1, 2 * x + 2 * M + 1);
        a[x + M].r++;
        b[x * 2 + 2 * M].r++;
        num[x * 3 + 3 * M]++;
    }
    int len = 1;
    while (len < len1) len <<= 1;
    fft(a, len, 1);
    fft(b, len, 1);
    for (int i = 0; i < len; i++) {
        a[i] = a[i] * a[i] * a[i] - a[i] * b[i] - a[i] * b[i] - a[i] * b[i];
    }
    fft(a, len, -1);
    for (int i = 0; i < len; i++) {
        ans[i] = (ll)(a[i].r + 0.5);
    }
    for (int i = 0; i < len; i++) {
        ll res = (ans[i] + 2 * num[i]) / 6;
        if (res > 0) {
            printf("%d : %lld\n", i - 3 * M, res);
        }
    }
    return 0;
}