Codeforces 1188E - Problem from Red Panda（找性质+组合数学）

Codeforces 题面传送门 & 洛谷题面传送门

咦，题解搬运人竟是我？

一道很毒的计数题。

先转化下题意，每一次操作我们可以视作选择一种颜色并将其出现次数 \(+k\)，之后将所有颜色的出现次数 \(-1\)。我们假设第 \(i\) 种颜色被操作了 \(c_i\) 次，那么一组 \(\{c_1,c_2,\cdots,c_k\}\) 符合条件当且仅当 \(\forall i,a_i+kc_i\ge\sum\limits_{i=1}^kc_i\)。我们所求即是符合这样的条件的 \(\{a_i-kc_i-\sum\limits_{i=1}^kc_i\}\) 的个数。

直接统计显然不行，因此考虑发掘一些性质。一个非常自然的猜想是，如果操作不能无限进行下去，那操作最多进行的轮数不会太多，大概就 \(\mathcal O(k)\) 级别的，因为如果存在一种操作序列满足 \(k\) 步之后仍然不会挂，那么我们一直重复这 \(k\) 次操作的过程中即可将游戏一直进行下去。因此从这个角度入手作文章。考虑对于一种颜色 \(i\)，如果我们希望操作能够继续下去，那么必然有前 \(a_i+1+ck\) 次操作中必须至少有 \(c+1\) 次操作作用在这个颜色上，因此我们考虑将数轴上这些形如 \(a_i+1+ck(c\ge 0)\) 的位置打上 \(+1\) 标记，然后对整个数轴进行一遍前缀和，我们假设得到的前缀和数组为 \(s_i\)，如果我们发现某个 \(s_i\) 大于 \(i\)，那么我们显然没办法安排这 \(i\) 次操作符合限制，也就表明操作次数最多为 \(i-1\)，break 掉即可。如果对于 \(i\in[1,k-1]\) 都不存在这样的情况则说明操作可以无限进行下去。

考虑怎样统计答案，首先是有限次操作的情况。需要注意到一个性质，那就是对于所有 \(x,y\in[1,k-1]\)，如果 \(x\ne y\)，那么所有操作 \(x\) 次后得到的序列肯定不同于操作 \(y\) 次后得到的序列，因为至少要 \(k\) 次操作可以将一个序列复原，而根据上面的推论，有限次操作的情况中操作次数的上界为 \(k-1\)，因此我们考虑枚举操作次数 \(x\)，那么我们考虑统计 \(x\) 次操作可以产生多少组不同的 \(\{c_1,c_2,\cdots,c_k\}\)。这个可以通过调用我们之前求得的前缀和数组 \(s_x\) 计算：有 \(s_x\) 次操作选择的颜色已经确定了，因此我们只能安排剩余 \(x-s_x\) 次操作选择的颜色，而这等价于 \(\sum\limits_{i=1}^kd_i=x-s_x\) 的非负整数解的组数，隔板法可算得方案数为 \(\dbinom{x-s_x+k-1}{k-1}\)。对于所有 \(x\) 计算一遍上式的值并将答案加起来即可。

接下来是无限次操作的情况。首先注意到一个性质，就是由于操作可以无限进行下去，对于任意 \(p\)，如果一个序列 \(\{a'\}\) 可以通过 \(p\) 次操作得到，那序列 \(\{a'\}\) 也可以通过 \(p+k\) 次操作得到。但这个结论反过来不一定成立，因为可能存在 \(p\) 过小而导致某些颜色无法操作的情况。不过这个问题比较容易解决，如果 \(p>\max{a_i}\) 就不会存在步数过小而无法操作全部颜色的情况了。因此直接对 \(x\in[10^6+1,10^6+k]\) 重复一遍上面的过程即可。

时间复杂度 \(\mathcal O(\max\{a_i\}+k)\)。

const int MAXN=1e6;
const int MOD=998244353;
int n,a[MAXN+5],cnt[MAXN*2+5],fac[MAXN*3+5],ifac[MAXN*3+5];
void init_fac(int n){
	for(int i=(fac[0]=ifac[0]=ifac[1]=1)+1;i<=n;i++) ifac[i]=1ll*ifac[MOD%i]*(MOD-MOD/i)%MOD;
	for(int i=1;i<=n;i++) fac[i]=1ll*fac[i-1]*i%MOD,ifac[i]=1ll*ifac[i-1]*ifac[i]%MOD;
}
int binom(int x,int y){
	if(x<0||y<0||x<y) return 0;
	return 1ll*fac[x]*ifac[y]%MOD*ifac[x-y]%MOD;
}
int main(){
	scanf("%d",&n);init_fac(MAXN*3);
	int lim=n+MAXN,res=0;
	for(int i=1;i<=n;i++){
		scanf("%d",&a[i]);
		for(int j=a[i]+1;j<=lim;j+=n) cnt[j]++;
	}
	for(int i=1;i<=lim;i++){
		cnt[i]+=cnt[i-1];
		if(cnt[i]>i){lim=i-1;break;}
	}
	if(lim<=MAXN){
		for(int i=0;i<=lim;i++) res=(res+binom(i-cnt[i]+n-1,n-1))%MOD;
	} else {
		for(int i=MAXN+1;i<=lim;i++) res=(res+binom(i-cnt[i]+n-1,n-1))%MOD;
	}
	printf("%d\n",res);
	return 0;
}