kmjp's blog

競技プログラミング参加記です

yukicoder : No.840 ほむほむほむら

こっちの方がとっつきやすいですね。
https://yukicoder.me/problems/no/840

問題

整数N,Kが与えられる。
「ほ」「む」「ら」で構成されるN文字の文字列Sのうち、S[i]="ほ"、S[j]="む"、S[k]="ら"となる(i<j<k)を満たす(i,j,k)のタプルの数がKの倍数となるようなものの数を求めよ。

解法

まず文字列が決まっている場合のタプルの数の数え方を考える。
これは定番で、

  • f(n) := n文字目までに登場した「ほ」の数
  • g(n) := n文字目までのうち2文字抜き出して作れる「ほむ」の数
  • k(n) := n文字目までのうち3文字抜き出して作れる「ほむら」の数

とすると、(n+1)文字目にほ・む・らのそれぞれが来る場合、以下のように加算される。

  • f(n+1) = f(n)+1
  • g(n+1) = g(n)+f(n)
  • h(n+1) = h(n)+g(n)

で計算できる。
これを踏まえて以下を考えよう。
P(n,i,j,k) := n文字目までについて、f(n)%K=i、g(n)%K=j、h(n)%K=kとなるような文字列の組み合わせ
P(N,*,*,0)が求めたい解となる。
i,j,kは合計でK^3通りなので、(K^3)次行列でPの漸化式を表現すれば、行列累乗で求められる。

int N,K;
ll mo=998244353;

const int MAT=125;
ll G[MAT][MAT],G2[MAT][MAT];

void powmat(ll p,int n,ll a[MAT][MAT],ll r[MAT][MAT]) {
	int i,x,y,z;
	ll a2[MAT][MAT];
	FOR(x,n) FOR(y,n) a2[x][y]=a[x][y];
	FOR(x,n) FOR(y,n) r[x][y]=(x==y);
	while(p) {
		ll h[MAT][MAT];
		if(p%2) {
			FOR(x,n) FOR(y,n) h[x][y]=0;
			FOR(x,n) FOR(z,n) FOR(y,n) h[x][y] += (r[x][z]*a2[z][y]) % mo;
			FOR(x,n) FOR(y,n) r[x][y]=h[x][y]%mo;
		}
		FOR(x,n) FOR(y,n) h[x][y]=0;
		FOR(x,n) FOR(z,n) FOR(y,n) h[x][y] += (a2[x][z]*a2[z][y]) % mo;
		FOR(x,n) FOR(y,n) a2[x][y]=h[x][y]%mo;
		p>>=1;
	}
}



void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,K) FOR(j,K) FOR(k,K) {
		G[(i+1)%K*K*K+j*K+k][i*K*K+j*K+k]++;
		G[i*K*K+(j+i)%K*K+k][i*K*K+j*K+k]++;
		G[i*K*K+j*K+(k+j)%K][i*K*K+j*K+k]++;
	}
	powmat(N,K*K*K,G,G2);
	ll ret=0;
	FOR(i,K) FOR(j,K) {
		ret+=G2[i*K*K+j*K][0];
	}
	cout<<ret%mo<<endl;
	
	
}

まとめ

これは★3か3.5でもよさそうな難易度に見えるな。