CF993E Nikita and Order Statistics

Link

Solution

考虑先将序列 \(<a_i>\) 转换，定义序列 \(<b_i>\)。令 \(b_i=[a_i<X]\)，那么对 \(b_i\) 做前缀和的话，得到的 \(s_i\) 就表示前 \(i\) 个数中小于 \(X\) 的数的个数。那么对于一个特定的 \(k\) ，我们就需要回答 \(s_i=s_j+k\) 的有序对 \((i,j)\) 的个数。我们成功地将区间计数转换成了元素统计，这实际上是这类问题的惯用伎俩。

再进一步，\(s\) 相同的前缀可以合并，令 \(c_i\) 表示 \(s=i\) 的前缀个数。那么对于答案 \(G(k)\) ，就有

\[G(k)=\sum_{i=0}^{m-k} c_i\times c_{i+k} \]

其中 \(m\) 为所有小于 \(X\) 的数的个数，容易知道当 \(i>m-k\) 时 \(c_{i+k}\) 都为零，所以不用统计。

到了这一步就很简单了，容易发现是一个差卷积的形式。具体地，令 \(A_{i}=c_i\) 和 \(B_{i}=c_{m-i}\)，那么上式就可以写成

\[G(k)=\sum_{i=0}^{m-k} A_{i}\times B_{m-i-k} \]

求出 \(A\) 和 \(B\) 的卷积后再把前 \(m\) 项 reverse 一下，就能得到 \(G\) 序列。复杂度 \(O(n\log n)\)。

特别的，\(G(0)\) 并不能这么处理，我们可以单独算出 \(G(0)\)，复杂度 \(O(n)\)。

\[G(0)=\sum_{i=0}^m \frac{c_i(c_i-1)}{2} \]

#include<stdio.h>
#include<math.h>
#include<algorithm>
using namespace std;
#define N ((1<<21)+3)
#define db double
#define ll long long

const db PI=acos(-1);

inline int read(){
    int x=0,flag=1; char c=getchar();
    while(c<'0'||c>'9'){if(c=='-') flag=0;c=getchar();}
    while(c>='0'&&c<='9'){x=(x<<1)+(x<<3)+c-48;c=getchar();}
    return flag? x:-x;
}

struct Complex{
    db x,y;
    Complex(db x_=0,db y_=0):x(x_),y(y_){}
    Complex operator +(const Complex X){return Complex(x+X.x,y+X.y);}
    Complex operator -(const Complex X){return Complex(x-X.x,y-X.y);}
    Complex operator *(const Complex X){return Complex(x*X.x-y*X.y,x*X.y+y*X.x);}
}a[N];

int n,X,m=0,op=0,rk[N];
ll G[N];

void FFT(){
    for(int i=0;i<n;i++)
        if(i<rk[i]) swap(a[i],a[rk[i]]);
    for(int p=2;p<=n;p<<=1){
        int len=p>>1;
        Complex w(cos(2*PI/p),op*sin(2*PI/p));
        for(int k=0;k<n;k+=p){
            Complex now(1.0,0.0);
            for(int l=k;l<k+len;l++){
                Complex t=now*a[l+len];
                a[l+len]=a[l]-t;
                a[l]=a[l]+t;
                now=now*w;
            }
        }
    }
}

int main(){
    n=read(),X=read();
    a[0]=1; int n_=n;ll ret=0;
    for(int i=1;i<=n;i++){
        int x=read();
        if(x<X) a[++m].x=1;
        else a[m].x++;
    }
    for(int i=0;i<=m;i++) a[m-i].y=a[i].x,ret+=(ll)a[i].x*(a[i].x-1)/2;
    for(n=1;n<=(m<<1);n<<=1);
    for(int i=0;i<n;i++)
        rk[i]=(rk[i>>1]>>1)|((i&1)? n>>1:0);
    op=1,FFT();
    for(int i=0;i<n;i++) a[i]=a[i]*a[i];
    op=-1,FFT();
    for(int i=0;i<=m;i++) G[i]=(ll)(a[i].y/n/2+0.5);
    reverse(G,G+1+m),G[0]=ret;
    for(int i=0;i<=n_;i++) printf("%lld ",G[i]);
}