51nod 算法马拉松 34 Problem D 区间求和2 （FFT加速卷积）

题目链接  51nod 算法马拉松 34  Problem D

在这个题中$2$这个质数比较特殊，所以我们先特判$2$的情况，然后仅考虑大于等于$3$的奇数即可。

首先考虑任意一个点对$(i, j)$，满足$1 <= i <= j <= n$

我们考虑这个点对对答案的贡献。

首先显然$i$和$j$必须有相同的奇偶性，那么$i + j$一定为偶数。

包含这个点对的有效的质数的区间长度为$[j - i + 1, min(i + j - 1, 2n + 1 - i - j)]$中的所有质数。

（就假设以这个区间为中心，一点点向外扩，直到左边界碰到$1$或者右边界碰到$n$位置）

设$f(x)$为前$x$个数中的质数个数（这里我们假设$2$不是质数），

那么这个点对对答案的贡献为$f_{min(i + j - 1, 2n + 1 - i - j)} * ∑a_{i} * a_{j} - f_{j - i} * ∑a_{i} * a_{j}$

现在分类讨论

• 当$i + j - 1 <= n$时，$min(i + j - 1, 2n + 1 - i - j) = i + j - 1$, 贡献为$f_{i + j - 1} * ∑a_{i} * a_{j} - f_{j - i} * ∑a_{i} * a_{j}$
• 当$i + j - 1  >  n$时，$min(i + j - 1, 2n + 1 - i - j) = 2n + 1 - i - j$， 贡献为$f_{2n + 1 - i - j} * ∑a_{i} * a_{j} - f_{j - i} * ∑a_{i} * a_{j}$

可以发现两种情况对于减掉的部分是一样的，那么合起来考虑。

对于那些贡献为正的部分，$∑a_{i} * b_{j}$用FFT加速卷积就可以求出；

对于那些贡献为负的部分，我们其实要求的就是$s_{j} * ∑a_{i} * a_{i - j}$，

令$b_{i} = a_{n - i + 1}$，那么问题转化成了求$a$和$b$的卷积和，问题也解决了。

时间复杂度$O(nlogn)$

我本来想用NTT，结果TLE了整整一天，准备放弃的时候发现FFT足够了，我真笨

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>

using namespace std;

#define rep(i, a, b)	for (int i(a); i <= (b); ++i)
#define	dec(i, a, b)	for (int i(a); i >= (b); --i)

typedef long long LL;

const double PI = acos(-1.0);
const int N = (1 << 18) + 10;
const LL mod = 1e9 + 7;

int s[N];
int n;
int cnt = 0;
LL ans;
LL a[N], b[N], c[N];

struct Complex{
	double x, y;
	Complex(double x = 0.0, double y = 0.0) : x(x), y(y){}
	Complex operator + (const Complex &b) const{
		return Complex(x + b.x, y + b.y);
	}
	Complex operator - (const Complex &b) const{
		return Complex(x - b.x, y - b.y);
	}
	Complex operator * (const Complex &b) const{
		return Complex(x * b.x - y * b.y, x * b.y + y * b.x);
	}
};

Complex x1[N << 1], x2[N << 1];

void change(Complex y[], int len){
	for (int i = 1, j = len / 2; i < len - 1; i++){
		if (i < j) swap(y[i], y[j]);
		int k = len / 2;
		while (j >= k){
			j -= k;
			k /= 2;
		}
		if (j < k) j += k;
	}
}

void fft(Complex y[], int len, int on){
	change(y, len);
	for (int h = 2; h <= len; h <<= 1){
		Complex wn(cos(-on * 2 * PI / h), sin(-on * 2 * PI / h));
		for (int j = 0; j < len; j += h){
			Complex w(1, 0);
			for (int k = j; k < j + h / 2; k++){
				Complex u = y[k];
				Complex t = w * y[k + h / 2];
				y[k] = u + t;
				y[k + h / 2] = u - t;
				w = w * wn;
			}
		}
	}

	if (on == -1){
		rep(i, 0, len - 1) y[i].x /= len;
	}
}

void mul(LL p[], int dp, LL q[], int dq){
	int len = 1;
	while (len <= dp + dq) len <<= 1;

	rep(i, 0, dp) x1[i] = Complex(p[i], 0);
	rep(i, dp + 1, len - 1) x1[i] = Complex(0, 0);

	rep(i, 0, dq) x2[i] = Complex(q[i], 0);
	rep(i, dq + 1, len - 1) x2[i] = Complex(0, 0);

	fft(x1, len, 1);
	fft(x2, len, 1);

	rep(i, 0, len - 1) x1[i] = x1[i] * x2[i];
	fft(x1, len, -1);
	rep(i, 0, dp + dq) c[i] = (LL)(x1[i].x + 0.5), c[i] %= mod;
}

int main(){

	rep(i, 3, (1 << 18)){
		int bl = sqrt(i), flag = 1;
		rep(j, 2, bl) if (i % j == 0){
			flag = 0;
			break;
		}

		s[i] = s[i - 1] + flag;
	}

	scanf("%d", &n);
	ans = 0;

	rep(i, 1, n){
		scanf("%lld", a + i);
		b[i] = a[i];
	}

	mul(a, n, b, n);

	rep(i, 1, n - 1){
		ans += a[i] * a[i + 1] % mod * 2 % mod;
		if (ans > mod) ans -= mod;
	}

	rep(i, 4, n + 1){
		if (i & 1) continue;
		ans = ans + c[i] * s[i - 1] % mod;
		if (ans > mod) ans -= mod;
	}

	rep(i, n + 2, n + n){
		if (i & 1) continue;
		ans = ans + c[i] * s[2 * n + 1 - i] % mod;
		if (ans > mod) ans -= mod;
	}

	reverse(b + 1, b + n + 1);
	mul(a, n, b, n);

	rep(i, 0, n - 1){
		if (i & 1) continue;
		ans = ans - c[i + n + 1] * s[i] % mod * 2 % mod;
		ans += mod;
		ans %= mod;
	}

	printf("%lld\n", ans);
	return 0;
}