【洛谷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;
}