[atAGC052C]Nondivisible Prefix Sums

当1为$a_{i}$中出现次数最多的元素（之一），则有以下结论——

结论：$a_{i}$合法当且仅当$P otmid sum_{i=1}^{n}a_{i}$且$sum_{i=1}^{n}[a_{i}=1]le (P-1)+sum_{1le ile n,a_{i} e 1}(P-a_{i})$

证明：

必要性——

若$Pmid sum_{i=1}^{n}a_{i}$，则取整体作为前缀即不满足条件

同时，当1的个数多于该值，假设有$k$个非1的数（即$sum_{i=1}^{n}[a_{i} e 1]$），则$sum_{i=1}^{n}a_{i}ge (k+1)P$

根据$P otmid sum_{i=1}^{n}a_{i}$，即$sum_{i=1}^{n}a_{i}>(k+1)P$

考虑求其前缀和，对于其中前缀和恰大于$P$、$2P$、……、$(k+1)P$的位置（不难证明这些位置必然存在且各不相同），若恰大于$tP$的位置上为1，同时其上一个位置必然小于等于$tP$，即恰好为$tP$，即不合法

因此，这些位置上必然都不为1，因此至少有$k+1$个非1的数，矛盾

充分性——

构造其重新排列后的$a'_{i}$，具体来说，当已经确定$a'_{1},a'_{2},...,a'_{k}$后，来确定$a'_{k+1}$

令$x$为$a_{i}$剩下（除去$a'_{1},a'_{2},...,a'_{k}$）的数中出现次数最多的元素（之一），对其分类讨论：

1.若$P otmid sum_{i=1}^{k}a'_{i}+x$，令$a'_{k+1}=x$

2.若$Pmid sum_{i=1}^{k}a'_{i}+x$，任取另一个剩下的元素$y e x$，令$a'_{k+1}=y$且$a'_{k+2}=x$

关于上述两者，唯一有可能导致其不合法的即第2类中不存在$y$的情况，即仅剩下的元素仅含$x$

若剩下的元素中仅有1个$x$，即与$P otmid sum_{i=1}^{n}a_{i}$矛盾，因此至少剩下两个$x$

在这种情况下，构造过程中，任意时刻出现次数最多（且唯一最多）的元素都是$x$

关于这个结论，考虑当某一时刻，某一个元素出现次数大于等于$x$出现次数-1，则简单分类讨论：

1.此时$x$作为出现次数最多的元素，其出现次数必然仍大于等于$x$出现次数-1

2.此时$x$不作为出现次数最多的元素，则假设为$y$，至多使用一个$y$，那么$y$出现次数大于等于$x$出现次数-1

换言之，由此可以说明之后任意时刻都存在一个元素出现次数大于等于$x$的出现次数-1

而最终，不存在这样的元素（其余元素出现次数都为0，而$x$出现次数至少为2），因此即任意时刻不存在出现次数大于等于$x$出现次数-1的元素，进而推出$x$任意时刻都是唯一最多的元素

由此，可以得到$x$为初始状态中出现次数最多的元素，即有$x=1$

根据构造过程，可以发现$a_{i}$之后恰好会有$P-a_{i}$个1，以及最初也必然有$P-1$个1，此时若还有剩下的1，即可得到$sum_{i=1}^{n}[a_{i}=1]>(P-1)+sum_{1le ile n,a_{i} e 1}(P-a_{i})$，矛盾

因此，即证明不存在这种不合法的情况，构造成立

而对于普通的序列$a_{i}$（即1不为$a_{i}$中出现次数最多的元素），若$x$为其中出现次数最多的元素，将所有元素在模$P$意义下除以$x$，得到序列$b_{i}equiv frac{a_{i}}{x}(mod P)$，不难证明$a_{i}$合法等价于$b_{i}$合法

同时，$b_{i}$合法根据结论即等价于$P otmid sum_{i=1}^{n}b_{i}$且$sum_{i=1}^{n}[b_{i}=1]le (P-1)+sum_{1le ile n,b_{i} e 1}(P-b_{i})$，同时前者又等价于$P otmid sum_{i=1}^{n}a_{i}$

接下来，考虑统计答案，答案可以通过容斥计算，即等于$P otmid sum_{i=1}^{n}a_{i}$的序列数，减去$P otmid sum_{i=1}^{n}a_{i}$且$sum_{i=1}^{n}[b_{i}=1]>(P-1)+sum_{1le ile n,b_{i} e 1}(P-b_{i})$的序列数

对于前者，记为$f_{n}$，则有
$$
f_{n}=(P-2)f_{n-1}+(P-1)((P-1)^{n-1}-f_{n-1})=(P-1)^{n}-f_{n-1}
$$
（初始状态为$f_{0}=0$）

根据等比数列求和，答案即
$$
f_{n}=(-1)^{n}sum_{i=1}^{n}(1-P)^{i}=(-1)^{n}frac{(1-P)-(1-P)^{n+1}}{P}=frac{(P-1)^{n+1}-(-1)^{n}(P-1)}{P}
$$
对于后者，根据式子，此时$a_{i}$中最大值出现必然唯一

另一方面，由于一个模$P$意义下的完系中所有数乘上$x$（满足$P otmid x$）后仍然是完系，即最大值并不影响每一个数出现的概率（情况），即可以统计1为出现次数最多的情况，并乘上$P-1$即可

（更通俗的来说，在$[1,P)$中随机一个数在乘上$x$并对$P$取模，等价于直接在$[1,P)$中随机一个数）

接下来，枚举1的个数$k$，即要求$sum_{1le ile n,a_{i} e 1}(P-a_{i})-k otmid P$且$sum_{1le ile n,a_{i} e 1}(P-a_{i})le k-P$

根据后者，这些数的和是$o(n)$的，即用$f_{i,j}$表示$i$个$[1,P-2]$之间的数，和为$j$的方案数，注意到两维都是$o(n)$级别的，且转移通过前缀和可以优化到$o(1)$，总复杂度即$o(n^{2})$

最终，（这一部分）答案即$sum_{i=0}^{n}{nchoose i}sum_{0le jle i-P,j otequiv i(mod P)}f_{n-i,j}$

总复杂度为$o(n^{2})$，可以通过

#include<bits/stdc++.h>
using namespace std;
#define N 5005
#define mod 998244353
int n,p,ans,c[N][N],f[N][N];
int pow(int n,int m){
    int s=n,ans=1;
    while (m){
        if (m&1)ans=1LL*ans*s%mod;
        s=1LL*s*s%mod;
        m>>=1;
    }
    return ans;
}
int main(){
    for(int i=0;i<N;i++){
        c[i][0]=c[i][i]=1;
        for(int j=1;j<i;j++)c[i][j]=(c[i-1][j-1]+c[i-1][j])%mod;
    }
    scanf("%d%d",&n,&p);
    ans=pow(p-1,n+1);
    if (n&1)ans=(ans+p-1)%mod;
    else ans=(ans+mod-p+1)%mod;
    ans=1LL*ans*pow(p,mod-2)%mod;
    f[0][0]=1;
    if ((n%p)&&(n>=p))ans=(ans+mod-p+1)%mod;
    for(int i=1;i<=n;i++){
        for(int j=1;j<=n;j++)f[i-1][j]=(f[i-1][j-1]+f[i-1][j])%mod;
        for(int j=1;j<=n;j++){
            if (j-(p-2)<=0)f[i][j]=f[i-1][j-1];
            else f[i][j]=(f[i-1][j-1]-f[i-1][j-(p-2)-1]+mod)%mod;
            if ((j<=n-i-p)&&(j%p!=(n-i)%p))ans=(ans+mod-1LL*(p-1)*c[n][i]%mod*f[i][j]%mod)%mod;
        }
    }
    printf("%d",ans);
}