Codeforces ECR #121 : F. A Random Code Problem

コードは短いんだけどね。
https://codeforces.com/contest/1626/problem/F

問題

N要素の整数列Aと、整数Kが与えられる。
以下のアルゴリズムを考える。

• i=1～Kまで以下をループする。
  • Aの要素vを1つランダムに選択し、解に加算する。その後、vからvをiで割った余りを引く。

要素の選択法はN^K通りある。
この時の解の総和を求めよ。

解法

1～KのLCMをLとする。
A[i]=aL+bと表現できるとき、各A[i]の解への寄与はa*f(b)+g(b)のような形になる。
よって、bを0～(L-1)までそれぞれにおいてf(b),g(b)をDPで求めておけば、A[i]の解への寄与を求められる。

int N;
ll A,X,Y,M;
const ll mo=998244353;
ll pat[720720][19];
ll addsum[720720][19];
ll sum[720720][19];

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

int K;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	int g=1;
	for(i=1;i<=16;i++) {
		g=g/__gcd(i,g)*i;
	}

	cin>>N>>A>>X>>Y>>K>>M;
	
	FOR(i,g) pat[i][K+1]=1;
	for(j=K;j>=1;j--) {
		FOR(i,g) {
			// not select
			pat[i][j]=pat[i][j+1]*(N-1)%mo;
			sum[i][j]=sum[i][j+1]*(N-1)%mo;
			addsum[i][j]=addsum[i][j+1]*(N-1)%mo;
			// select
			k=i-i%j;
			(pat[i][j]+=pat[k][j+1])%=mo;
			(sum[i][j]+=(sum[k][j+1]+pat[k][j+1]))%=mo;
			(addsum[i][j]+=(addsum[k][j+1]+i*pat[k][j+1]))%=mo;
		}
	}
	
	ll ret=0;
	FOR(i,N) {
		x=A%g;
		(ret+=A/g*g*sum[x][1]+addsum[x][1])%=mo;
		A=(A*X+Y)%M;
	}
	cout<<ret<<endl;
	
}

まとめ

気付いてしまえばすぐなんだけどね。