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

时间:2022-08-02 18:45:44

题目链接  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;
}