CodeForces - 993E: Nikita and Order Statistics (FFT)
Nikita likes tasks on order statistics, for example, he can easily find the kk-th number in increasing order on a segment of an array. But now Nikita wonders how many segments of an array there are such that a given number xx is the kk-th number in increasing order on this segment. In other words, you should find the number of segments of a given array such that there are exactly kk numbers of this segment which are less than xx.
Nikita wants to get answer for this question for each kk from 00 to nn, where nn is the size of the array.
Input
The first line contains two integers nn and xx (1≤n≤2⋅105,−109≤x≤109)(1≤n≤2⋅105,−109≤x≤109).
The second line contains nn integers a1,a2,…,ana1,a2,…,an (−109≤ai≤109)(−109≤ai≤109) — the given array.
Output
Print n+1n+1 integers, where the ii-th number is the answer for Nikita's question for k=i−1k=i−1.
Examples
5 3
1 2 3 4 5
6 5 4 0 0 0
2 6
-5 9
1 2 0
6 99
-1 -1 -1 -1 -1 -1
0 6 5 4 3 2 1
来自luogu的题解,解释得很清楚,我就不多说了。
解题思路
显然将所有小于 x 的数转化为1,其他的为0,这就是个01序列,然后求区间和是 k 的区间个数。
那么我们首先要取一个左端点,再取一个右端点。
k要取遍,并且是 10^5 级的数据,考虑FFT。
考虑左端点取每一个点时,在这个点左边的 11 的个数,若个数为 tt ,则 x^t的系数加1,这样构造一个多项式 A 。
右端点取每一个点时,在这个点右边的 11 的个数,构造一个类似的多项式 B 。
那么 A 和 B 做卷积,第 t 项的系数就是区间外 1 的个数为 t 的区间个数。
桥豆麻袋,有可能出现右端点取得比左端点小的情况。
发现这种情况下,算出来在区间外 1 的个数一定大于等于整个序列 1 的个数,所以受影响的只有 k=0 的情况,所以单独用组合的方法计算一下 k=0 的情况即可。
(懒得写虚数和重载的可以直接套C++的complex,虽然会稍微慢一点,但是代码简洁。。。
#include<bits/stdc++.h> #define ll long long using namespace std; typedef complex<double> cp; const int N=530000; const double pi=acos(-1); cp a[N],b[N],W,w,p; int n=1,m,i,j,k,x,sum,R[N],lst=-1;ll ans; inline void fft(cp*c,int t) { for(i=0;i<n;i++) R[i]<i?swap(c[R[i]],c[i]),0:0; for(i=1;i<n;i<<=1) for(j=0,W={cos(pi/i),sin(pi/i)*t};j<n;j+=i<<1) for(k=0,w={1,0};k<i;k++,w=w*W) p=c[j+k+i]*w,c[j+k+i]=c[j+k]-p,c[j+k]=c[j+k]+p; } int main() { scanf("%d%d",&m,&k); for(i=0;i<m;i++){ a[sum]+=cp(1,0),scanf("%d",&x); x=x<k; sum+=x,b[sum]+=cp(1,0); if(x) lst=i; ans+=i-lst; } printf("%I64d ",ans); for(i=0;i<=sum>>1;i++) swap(a[i],a[sum-i]); for(;n<=sum<<1;n<<=1); for(i=1;i<n;i++) R[i]=R[i>>1]>>1|(i&1?n>>1:0); fft(a,1),fft(b,1); for(i=0;i<n;i++) a[i]=a[i]*b[i]; fft(a,-1); for(i=1;i<=m;i++) printf("%I64d ",(ll)((a[i+sum].real()+0.5)/n)); }