【洛谷P5641】【CSGRound2】开拓者的卓识

推式子+NTT

题目大意:

给定数列 $a_1,a_2,\dots,a_n$,定义以下函数:

给定 $k$,对所有 $p\in[1,n]$,求 $f(1,p,k)$ 对 $998,244,353$ 取模后的值。

题解:

考虑一个数 $a_i$ 在 $p=r$ 时的贡献。

相当于在 $1\sim i$ 中选 $k-1$ 个可以重复的位置,依次为左端点,$i\sim r$ 中选 $k-1$ 个可以重复的位置,依次作为右端点。

那么其贡献为 $a_i\times\binom{i+k-2}{k-1}\times \binom{r-i+k-1}{k-1}$

其上标之和是 $r+2k-3$,是个只与 $r$ 有关的数。

可以构造卷积,令 $A_i=a_i\times \binom{i+k-2}{k-1},B_i=\binom{i+k-1}{k-1}$。

然后答案就是 $A$ 与 $B$ 的卷积的第 $1\sim n$ 位。

时间复杂度 $O(n\log n)$。

Code:

#include<iostream>
#include<algorithm>
using namespace std;
const int N=262144,md=998244353,g3=(md+1)/3;
typedef long long LL;
int A[N],B[N],n,k,a[N],rev[N],lim;
inline void upd(int&a){a+=a>>31&md;}
inline int pow(int a,int b){
    int ret=1;
    for(;b;b>>=1,a=(LL)a*a%md)
        if(b&1)ret=(LL)ret*a%md;
    return ret;
}
inline void init(int n){
    int l=-1;
    for(lim=1;lim<n;lim<<=1)++l;
    for(int i=1;i<lim;++i)
        rev[i]=(rev[i>>1]>>1)|((i&1)<<l);
}
void NTT(int*a,int f){
    for(int i=1;i<lim;++i)
        if(i<rev[i])swap(a[i],a[rev[i]]);
    for(int i=1;i<lim;i<<=1){
        const int gi=pow(f?3:g3,(md-1)/(i<<1));
        for(int j=0;j<lim;j+=i<<1)
            for(int k=0,g=1;k<i;++k,g=(LL)g*gi%md){
                const int x=a[j+k],y=(LL)g*a[j+k+i]%md;
                upd(a[j+k]+=y-md),upd(a[j+k+i]=x-y);
            }
    }
    if(!f){
        const int iv=pow(lim,md-2);
        for(int i=0;i<lim;++i)a[i]=(LL)a[i]*iv%md;
    }
}
int main(){
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    cin>>n>>k;
    for(int i=1;i<=n;++i)
        cin>>a[i];
    ++n;
    A[1]=1;
    for(int i=2;i<n;++i)A[i]=(LL)A[i-1]*pow(i-1,md-2)%md*(i+k-2)%md;
    for(int i=0;i<n-1;++i)B[i]=A[i+1];
    for(int i=1;i<n;++i)A[i]=(LL)A[i]*a[i]%md;
    init(n<<1);
    NTT(A,1),NTT(B,1);
    for(int i=0;i<lim;++i)A[i]=(LL)A[i]*B[i]%md;
    NTT(A,0);
    for(int i=1;i<n;++i)
        cout<<A[i]<<' ';
    return 0;
}