CF993E Nikita and Order Statistics

考虑记录一个前缀和(s),当(a_i < x)时这一位为(1)，否则为(0);
记录一下(s[i])的数量(f[i])
那么就有(ans_k = sum_{i = 0}^{n - k}f_i * f_{i + k})
FFT经典操作，翻转数组。
设(g[i] = f[n - i])
那么发现(ans_i = h[n - i] = sum f[j]g[n - i - j])
于是计算卷积就好了，不过注意数据范围和(ans_0)要特判。

CF993E Nikita and Order Statistics

#include<iostream>
#include<cstdio>
#include<complex>
#define ll long long
#define comp std::complex<double>
#define fft fast_fast_tle
#define int long long

const int N = (1 << 20) + 10 << 1;
const double PI = std::acos(-1.0);

inline ll read(){
	ll ans = 0,f = 1;
	char a = getchar();
	while(!(a <= '9' && a >= '0') && (a != '-'))a = getchar();
	if(a == '-')
	f = -1,a = getchar(); 
	while(a <= '9' && a >= '0')
	ans = (ans << 3) + (ans << 1) + (a & 15),a = getchar();
	return ans * f;
}

ll n,m,lim,r[N];
comp a[N],b[N];

void fft(comp *a,int type){
	for(int i = 0;i < lim;++i)
	if(i < r[i])
	std::swap(a[i],a[r[i]]);
	for(int i = 1;i < lim;i <<= 1){
		comp x(cos(PI / i),type * sin(PI / i));
		for(int j = 0;j < lim;j += (i << 1)){
			comp y(1,0);
			for(int k = 0;k < i;++k,y *= x){
			comp p = a[j + k],q = y * a[j + k + i];
			a[j + k] = p + q;a[j + k + i] = p - q;
			}
		}
	}
}
ll ni,k;
ll num[N];
ll s[N],f[N];

inline ll get0(){
	ll ans = 0;
	for(int i = 0;i <= ni;++i)
	ans += f[i] * (f[i] - 1) / 2;
	return ans;
}

signed main(){
	ni = read(),k = read();
	for(int i = 1;i <= ni;++i){
		s[i] = s[i - 1] + (read() < k);
//		std::cout<<s[i]<<std::endl;
		f[s[i]] ++ ; 
	}
	f[0] ++ ;
	n = ni,m = ni;
	for(int i = 0;i <= n;++i)
	a[i] = f[i];
	for(int i = 0;i <= m;++i)
	b[i] = f[ni - i];
	int l = 0;
	for(lim = 1;lim <= n + m;lim <<= 1)++l;
	for(int i = 0;i < lim;++i)
	r[i] = (r[i >> 1] >> 1) | ((i & 1) << (l - 1));
	fft(a,1),fft(b,1);
	for(int i = 0;i <= lim;++i)
	a[i] *= b[i];
	fft(a,-1);
	std::cout<<get0()<<" "; 
	for(int i = n - 1;i >= 0;--i)
	std::cout<<(int)(0.5 + a[i].real() / lim)<<" ";
}