题解[CF1628F]A_Random_Code_Problem

题意

给定一个数组 \(a\)，进行 \(k\) 次操作。第 \(i\) 操作等概率随机 \(a\) 中一个元素 \(a_x\)，将这个元素的值加入答案，并使其减去 \(a_x\bmod i\) 。问期望乘上 \(n^k\) 的值，对 \(998244353\) 取模。

\(n\le 10^7\)，\(k\le17\)，\(a_i\le998244353\)。

思路

期望乘上 \(n^k\) 就是所有情况下答案的和。

研究一下每个元素对答案的贡献，发现只和该元素被选中的次数集合有关，与其他元素无关。那么我们分开计算每个元素的贡献。

一个显然的做法是枚举每一次是否选中这个元素，并且计算贡献与方案数。这个 \(O(n2^k)\) 暴力较简单，不多讲。

注意到一个元素对答案的贡献与位置无关，只与值的大小有关，所以可以计算出每一个值对答案的贡献。记 \(f_{i,j}\) 表示有多少种方式使得某元素 \(i\) 次操作后大小为 \(j\)。注意，这里的 DP 相当于对每个元素进行 DP 的过程合并成一个 DP，所以 \(f_{i,j}\) 是所有元素 \(i\) 次操作后变成 \(j\) 的方案数之和。转移方程为 \(f_{i+1,j-j\bmod i}+=f_{i,j}\)，\(f_{i+1,j}+=(n-1)f_{i,j}\) 最后对于每个 \(f_{i,j}\) 计算贡献。

如何优化？这里有一个很妙的做法：观察性质，每次操作对元素的影响是减去 \(a_j\bmod i\)，这样的影响，最多也不会减去超过 \(a_j\bmod \mbox{lcm}_{i=1}^ki\)。为什么？因为任意 \(1\le i\le k\) 都是 \(\mbox{lcm}_{i=1}^ki\) 的因数，所以无论何时都有 \(a_j\bmod i\le a_j\bmod \mbox{lcm}_{i=1}^ki\)。那么我们只需要 DP \(a_j\bmod \mbox{lcm}_{i=1}^ki\) 的部分，而其他部分是不会随操作改变的，直接 \(O(n)\) 计算贡献即可。总复杂度 \(O(n+k\mbox{lcm}_{i=1}^ki)\)。

实现

直接 DP 的话空间开不下，可以选择滚动数组，也可以根据转移的方向开单个数组，同时卡时空常数。具体见代码。

inline void Calc(int x){
	MAdd(ans,Mul(x/l*l,Mul(m,pn[m-1])));
	MAdd(f[x%l],1);
}

int main(){
	l=1;
	n=Read();
	a[1]=Read();
	int x=Read(),y=Read();
	m=Read();
	int mod=Read();
	pn[0]=1;
	for(int i=1;i<m;++i)
		l=l/Gcd(l,i)*i,pn[i]=Mul(pn[i-1],n);
	Calc(a[1]);
	for(int i=2;i<=n;++i)
		a[i]=(1ll*a[i-1]*x+y)%mod,Calc(a[i]);
	for(int i=1;i<=m;++i)
		for(int j=0;j<l;++j){
			MAdd(ans,Mul(Mul(f[j],j),pn[m-i]));
			if(i<m){
				if(j%i)
					MAdd(f[j-j%i],f[j]),f[j]=Mul(f[j],n-1);
				else
					f[j]=Mul(f[j],n);
			}
		}
	printf("%d\n",ans);
	return 0;
}