AGC052C - Nondivisible Prefix Sums

对于所有值为\([1,p-1]\)中的整数的，长度为\(n\)的排列。统计：存在方案使重排列之后，所有前缀和不是\(p\)的倍数的排列个数。

\(n,p\le 5000\)

\(p\)为质数。

---

对于合法的集合，充要条件：满足总和不为\(p\)的倍数，设\(x\)为众数（出现\(c\)次），将所有数乘\(x^{-1}\)，设除了\(1\)之外的数为\(B_1,B_2,\dots,B_k\)，满足\(c\le \sum (p-B_i)+p-1\)。

证明：

必要性：如果不满足，则\(c\ge \sum(p-B_i)+p\)，又因为不满足\(\sum B_i+c\equiv 0\pmod p\)，所以\(\sum B_i+c\)最小值为\(p(k+1)+1\)。把前缀和视作在数轴上跳，当跳过\(p\)的倍数时，一定不能用\(1\)，然而剩下的数只有\(k\)个，不能跳过\(k+1\)次。得证。

充分性：考虑构造，设当前前缀和为\(s\)，每次取出剩下最多的数\(x\)，如果\((s+x)\mod p\neq 0\)，则直接丢到末尾；否则任选另一个数\(y\)，把\(y,x\)丢到末尾。如果出现无解，一定是：\((s+x)\mod p=0\)并且没有其它数，由于总和不是\(p\)的倍数，所以至少有两个\(x\)。发现在这个过程中，初始众数在集合中的出现次数永远不会小于当前众数在集合中的出现次数减一。所以如果抵达了无解的状况，剩下的至少两个一定是\(1\)。最坏情况下，前面应该是：\(p-1\)个\(1\)，\(B_1\)，\(p-B_1\)个\(1\)，\(B_2\)，……，\(B_k\)，\(p-B_k\)个\(1\)。然而由条件得此时\(1\)一定就耗尽了。

先算总和不为\(p\)的倍数的方案。设\(g_n\)，则有\(g_n=g_{n-1}(p-2)+((p-1)^{n-1}-g_{n-1})(p-1)\)。（分别考虑长度\(n-1\)的前缀是否总和为\(p\))

在总和中剔除不合法的，不合法时众数一定唯一，把它当\(1\)算最后乘上\(p-1\)。

设\(f_{cnt,sum}\)表示不是\(1\)的数用了\(cnt\)个，\(\sum p-B_i\)的和为\(sum\)的序列的个数。（注意此时不需要把像算可重排列那样把出现次数的阶乘的倒数乘起来，只需要直接往后面加变成计算不含\(1\)的排列）列出方程之后\(O(n^3)\)，可以简单优化至\(O(n^2)\)。算出来之后，对于要算到的\((cnt,sum)\)，满足\(n-cnt\ge sum-p\)和\((n-cnt-sum)\mod p\neq 0\)，算出把\(1\)插进去之后的方案加入答案。

总时间\(O(n^2)\)。

---

using namespace std;
#include <bits/stdc++.h>
#define N 5005
#define mo 998244353
#define ll long long
ll qpow(ll x,ll y=mo-2){
	ll r=1;
	for (;y;y>>=1,x=x*x%mo)
		if (y&1)
			r=r*x%mo;
	return r;
}
int n,p;
ll C[N][N];
ll f[N][N],g[N];
int main(){
	freopen("in.txt","r",stdin);
	scanf("%d%d",&n,&p);
	for (int i=0;i<=n;++i){
		C[i][0]=1;
		for (int j=1;j<=i;++j)
			C[i][j]=(C[i-1][j-1]+C[i-1][j])%mo;
	}
	f[0][0]=1;
	for (int i=1;i<=n;++i){
		ll s=0;
		for (int j=1;j<=n;++j){
			s+=f[i-1][j-1];
			if (j-(p-2)-1>=0)
				s-=f[i-1][j-(p-2)-1];
			f[i][j]=s%mo;
//			for (int k=1;k<=p-2 && j-k>=0;++k)
//				(f[i][j]+=f[i-1][j-k])%=mo;
		}
	}
	g[0]=0;
	for (int i=1;i<=n;++i)	
		g[i]=(qpow(p-1,i)-g[i-1])%mo;
//	for (int i=1;i<=n;++i)	
//		g[i]=(g[i-1]*(p-2)+(qpow(p-1,i-1)-g[i-1])*(p-1))%mo;
	ll sum=0;
	for (int c=0;c<=n;++c)
		for (int s=0;s<=n;++s)
			if (n-c>=s+p && (n-c-s)%p)
				(sum+=f[c][s]*C[n][c])%=mo;
	ll ans=(g[n]-sum*(p-1))%mo;
	ans=(ans+mo)%mo;
	printf("%lld\n",ans);
	return 0;
}