样例过了,但是交上去WA,自测没查出错
#define _c const
#define _r register
#define ALL(x) x.begin(), x.end()
_c double pi = acos(-1);
using namespace std;
struct complex_ {
double Real, Imagine;
complex_ (double Real = 0.0, double Imagine = 0.0) : Real(Real), Imagine(Imagine) {}
~complex_ () {}
friend complex_ operator+(_c complex_ &x, _c complex_ &y) { return complex_ (x.Real + y.Real, x.Imagine + y.Imagine); }
friend complex_ operator-(_c complex_ &x, _c complex_ &y) { return complex_ (x.Real - y.Real, x.Imagine - y.Imagine); }
friend complex_ operator*(_c complex_ &x, _c complex_ &y) { return complex_ (x.Real * y.Real - x.Imagine * y.Imagine, x.Real * y.Imagine + x.Imagine * y.Real); }
};
int32_t n;
_c size_t MAX_SIZE_log2 = 17;
_c size_t MAX_SIZE = 1 << 17;
size_t rev[MAX_SIZE];
complex_ F[MAX_SIZE], G[MAX_SIZE], H[MAX_SIZE];
int64_t ans[MAX_SIZE];
/*
容斥原理
Ans = (F^3(x) - 3G(x)F(x) + 2H(x)) / 6 的系数
*/
inline void FFT(complex_ *, size_t, int32_t);
inline void solve() {
FFT(F, MAX_SIZE, 1);
FFT(G, MAX_SIZE, 1);
FFT(H, MAX_SIZE, 1);
for (size_t i = 0; i < MAX_SIZE; i++) {
F[i] = F[i] * (F[i] * F[i] - complex_ (3, 0) * G[i]) + complex_(2, 0) * H[i];
}
FFT(F, MAX_SIZE, -1);
FFT(G, MAX_SIZE, -1);
FFT(H, MAX_SIZE, -1);
for (size_t i = 0; i < MAX_SIZE; i++) {
ans[i] = F[i].Real / 6.0 + 0.5;
if (ans[i]) {
printf("%u : %lld\n", i - 60015, ans[i]);
}
}
}
signed main() {
for (size_t i = 0; i < MAX_SIZE; i++) {
rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (MAX_SIZE_log2 - 1));
}
cin >> n;
while (n--) {
size_t inp;
cin >> inp;
F[inp + 20005].Real++;
G[(inp + 20005) * 2].Real++;
H[(inp + 20005) * 3].Real++;
}
solve();
return 0;
}
inline void FFT(complex_ *func, size_t size, int32_t opt) {
for (size_t i = 0; i < size; i++)
if (i < rev[i]) { swap(func[i], func[rev[i]]); }
for (size_t mid = 1; mid < size; mid <<= 1) {
complex_ timer(cos(opt * pi / mid), sin(opt * pi / mid));
for (size_t add = mid << 1, j = 0; j < size; j += add) {
complex_ num(1.0, 0.0);
for (size_t k = 0; k < mid; k++, num = num * timer) {
complex_ a = func[j + k], b = num * func[j + k + mid];
func[j + k] = a + b;
func[j + k + mid] = a - b;
}
}
}
if (opt == -1)
for (size_t i = 0; i < size; i++) {
func[i].Real /= double(size);
func[i].Imagine /= double(size);
}
}