题目描述
题目大意:一个非负整数组成的数列,记其前缀和为S。对于一段子区间,它对其部分和的贡献为其区间长度。问对于从0到Sn的部分和,其得到的总价值。
0 <=S<=50000
1<=n<=100000
T=5
思路
根据贡献放系数,位置为指数的套路,我们要想要构造出的答案为:
(i−j+1)si−sj−1
我们要构造生成函数来算出它,反正我是想不到的了,需要构造4个生成函数,然后将它们两两相乘(卷积),再相减,就像这样:
(∑ixsi)(∑x−si−1)−(∑xsi)(∑(i−1)x−si−1)
其中
1≤i≤n
,相乘与相减的皆是多项式。怎么样,聪明的你想到了吗?
然后FFT大力卷积即可,我们对负指数加上一个偏移量(Sn),由于无负数,所以算出来的1~Sn的结果不重不漏,但是0的话会算错(目测会少),因为i到j的区间的贡献正反抵消了,然后…不管了!
总之0要O(n)单独算,这里我居然算错了几次!?
代码
#include <bits/stdc++.h>
#define maxn 100010
using namespace std;
typedef long long LL;
typedef long double LD;
const LD PI = acos(-1.0);
int T, n, N, a[maxn], S[maxn], rev[maxn<<2];
struct Complex{
LD real, image;
Complex() {}
Complex(LD _real, LD _image){
real = _real; image = _image;
}
friend Complex operator + (Complex A, Complex B){
return Complex(A.real + B.real, A.image + B.image);
}
friend Complex operator - (Complex A, Complex B){
return Complex(A.real - B.real, A.image - B.image);
}
friend Complex operator * (Complex A, Complex B){
return Complex(A.real * B.real - A.image * B.image, A.image * B.real + A.real * B.image);
}
}f[maxn<<2], g[maxn<<2], h[maxn<<2];
void Init(int x){
int L = 0;
for(N = 1; N <= x; N <<= 1, L++);
for(int i = 0; i < N; i++) rev[i] = (rev[i>>1]>>1) | ((i&1)<<(L-1));
}
void FFT(Complex *A, int DFT){
for(int i = 0; i < N; i++) if(i < rev[i]) swap(A[i], A[rev[i]]);
for(int s = 1; (1<<s) <= N; s++){
int m = 1 << s;
Complex wn(cos(DFT*2*PI/m), sin(DFT*2*PI/m));
for(int k = 0; k < N; k += m){
Complex w(1, 0);
for(int j = 0; j < (m>>1); j++){
Complex u = A[k+j], t = w * A[k+j+(m>>1)];
A[k+j] = u + t;
A[k+j+(m>>1)] = u - t;
w = w * wn;
}
}
}
if(DFT == -1) for(int i = 0; i < N; i++) A[i].real /= N;
}
int main(){
scanf("%d", &T);
while(T --){
scanf("%d", &n);
for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
S[0] = 0;
for(int i = 1; i <= n; i++) S[i] = S[i-1] + a[i];
Init(S[n] + S[n]);
for(int i = 0; i < N; i++) f[i] = g[i] = h[i] = Complex(0, 0);
for(int i = 1; i <= n; i++){
f[S[i]].real += 1.0 * i;
g[S[n]-S[i-1]].real += 1.0;
}
FFT(f, 1);
FFT(g, 1);
for(int i = 0; i < N; i++) f[i] = f[i] * g[i], g[i] = Complex(0, 0);
FFT(f, -1);
for(int i = 1; i <= n; i++){
g[S[i]].real += 1.0;
h[S[n]-S[i-1]].real += 1.0 * (i-1);
}
FFT(g, 1);
FFT(h, 1);
for(int i = 0; i < N; i++) g[i] = g[i] * h[i];
FFT(g, -1);
int cnt = 0;
LL sum = 0;
for(int i = 1; i <= n; i++){
if(!a[i]) cnt ++;
if(a[i] || i == n){
sum += 1LL * cnt * (cnt+1) * (cnt+2) / 6;
cnt = 0;
}
}
for(int i = S[n]; i <= S[n]+S[n]; i++){
if(i == S[n]) printf("%I64d\n", sum);
else printf("%I64d\n", (LL)(f[i].real - g[i].real + .5));
}
}
return 0;
}