yukicoder : No.1887 K Consecutive Ks (Easy)

やることはシンプル。
https://yukicoder.me/problems/no/1887

問題

1以上M以下の整数からなるN要素の整数列Aのうち、以下を満たすものは何通りか。

• ある整数kがk個連続しているような区間が1つ以上ある。

解法

M^Nから、条件を満たさないものを引くことを考える。

dp(n,v) := 条件を満たさないn要素の整数列で、末尾がvのものの組み合わせ数
とする。
sum(n) := dp(n,v)のvに対する総和
とする。sum(n)・dp(n,v)がわかっている状況で、n+1要素目以降値wを続けようとする場合、wの直前の組み合わせは(sum(n)-dp(n,w)である。
また、wは(w-1)個以下まで連続できるため、以下のように解を計上できる。

• dp(n+1,w) += sum(n)-dp(n,w)
• dp(n+2,w) += sum(n)-dp(n,w)
• ...
• dp(n+w-1,w) += sum(n)-dp(n,w)

これを逐一行うと全体でO(NM^2)かかりTLEするが、上の式は足しこむ値が同じなので容易に累積和に持ち込んでO(NM)にできる。

int N,M;
const ll mo=998244353;

ll dp[6030][3030];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	
	for(i=1;i<=M;i++) {
		dp[1][i]=1;
		dp[i][i]+=mo-1;
	}
	for(i=1;i<=N;i++) {
		ll sum=0;
		for(j=1;j<=M;j++) {
			(dp[i][j]+=dp[i-1][j])%=mo;
			sum+=dp[i][j];
		}
		sum%=mo;
		for(j=1;j<=M;j++) {
			ll add=(sum+mo-dp[i][j])%mo;
			(dp[i+1][j]+=add)%=mo;
			(dp[i+j][j]+=mo-add)%=mo;
		}
	}
	
	ll ret=1;
	FOR(i,N) ret=ret*M%mo;
	FOR(i,M) ret+=mo-dp[N][i+1];
	cout<<ret%mo<<endl;
}

まとめ

Hardの方はまだ解いてないけどね。