八校联测 序列 (生成函数)

Posted 2021-05-07 nlkog

Description

你有一个长为n的序列{(a_n)}，每个位置你可以填一个[0,m- 1]中的整数

我们记{(a_n)}的前缀和为{(s_n)}，即:(s_i=sumlimits_{j=1}^{i}a_j)

问有多少个不同的序列{$ a_n $}满足至少有k个s;是m的倍数。

答案可能很大，请输出答案对998244353取模的结果。

Task

Input

第一行一个整数T表示数据组数。

以下T行，每行三个整数分别表示n,m,k。

Output

对于每组数据，输出一行一个整数表示答案对998244353取模的结果。

Sample Input

2
2 2 3
3 2 2

Sample Output

0
4

Solution

考虑dp

设f(i,j)表示在前 i 个数组成的序列中, 有且仅有 j 个(s_i)是m的倍数

因为(a_i)取值范围是[0,m-1], 所以对于(a_{i-1}), $a_i (共有m-1种取值使得)a_{i-1}+a_i$不是m的倍数

所以
[ f(i,j)=f(i-1,j-1)+f(i-1,j)*(m-1) ]
答案就是(sumlimits_{i=k}^n f(n,i))

时间复杂度O((n^2))

正确性get, 复杂度不够

考虑生成函数优化

设多项式(f_0,f_1,f_2...,f_n)

满足(f_i=sumlimits_{j=0}^n f(i,j)x^j)

假设(f_i=f_{i-1} imes g)

所以
[ f_i=f_{i-1} imes g=sum_{j=0}^n sum_{k=0}^j f_{i-1}(k) imes g(j-k)=sum_{j=0}^nf_{i-1}(j-1)+f_{i-1}(j) imes(m-1) /*结合dp式考虑*/ ]
解得
[ g(0)=m-1g(1)=1即g=m-1+x ]
因为(f_0=1)

所以
[ f_n=g^n=(m-1+x)^n=sum_{i=0}^nC_n^i (m-1)^{n-i}x^i/*二项式定理*/ ]
答案就是
[ ans=sum_{i=k}^nf_n(i)=sum_{i=k}^nC_n^i(m-1)^{n-i} ]

Code

#include<bits/stdc++.h>
#define int long long
using namespace std;
int read(){
    int x=0,f=1;char ch=getchar();
    for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-1;
    for(;isdigit(ch);ch=getchar())x=x*10+ch-'0';
    return x*f;
}
const int N=1e5+28,p=998244353;
int mul[N],inv[N],pw[N];
void Pre(int n=1e5){
    mul[0]=mul[1]=inv[1]=pw[0]=1;
    for(int i=2;i<=n;i++){
    mul[i]=mul[i-1]*i%p;
    inv[i]=(p-p/i)*inv[p%i]%p;
    }
    for(int i=2;i<=n;i++)inv[i]=inv[i]*inv[i-1]%p;
}
int C(int n,int m){
    if(n==m)return 1;
    return mul[n]*inv[m]%p*inv[n-m]%p;
}
signed main(){
    //freopen("sequence.in","r",stdin);
    //freopen("sequence.out","w",stdout);
    Pre();
    int t=read();
    while(t--){
    int n=read(),m=read(),k=read(),ans=0;
    pw[0]=1;
    for(int i=1;i<=n;i++){
        pw[i]=1ll*pw[i-1]*(m-1)%p;
    }
    for(int i=k;i<=n;i++){
        int tmp=1ll*C(n,i)*pw[n-i]%p;
        ans=1ll*(ans+tmp)%p;
    }
    printf("%lld
",ans);
    }
    return 0;
}
/*
2
2 2 3
3 2 2
*/