Codeforces #934 : Div1 D2. Counting Is Fun (Hard Version)

部分点は取れたが。
https://codeforces.com/contest/1943/problem/D2

問題

非負整数列Bに対し、以下を満たすものを良い整数と呼ぶ。

• Bのうち2以上の長さの区間を選び、区間内の要素をデクリメントする
• その作業を繰り返し、Bの全要素を0にする。

整数N,Kが与えられる。
0～Kの整数値を取るN要素の整数列Aは(K+1)^N通りある。
このうち良い数列は何通りか。

解法

まずO(N^3)解法を考える。
条件を満たすのは、A[i]＞A[i-1]+A[i+1]となるものがないことである。
末尾2要素を覚えながら先頭から値を決めて行くと、O(N^4)のDPとなり、累積和を使うとO(N^3)となる。

これを包除原理で状態を減らす。
先頭から要素を決めるのは同じだが、状態として末尾の値と、違反した箇所の個数の偶奇をもって計算していく。

int T,N,K;
ll mo;

ll dp[4040][4040][2];


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K>>mo;
		
		FOR(x,N+3) FOR(y,K+1) dp[x][y][0]=dp[x][y][1]=0;
		dp[0][0][0]=1;
		
		FOR(i,N+1) {
			ll s0=0,s1=0;
			FOR(x,K+1) {
				s0+=dp[i][x][0];
				s1+=dp[i][x][1];
			}
			s0%=mo;
			s1%=mo;
			
			ll ssa0=0,ssa1=0;
			ll ssb0=0,ssb1=0;
			if(i) {
				FOR(x,K+1) {
					(ssa0+=dp[i-1][x][0]*(K-x))%=mo;
					(ssa1+=dp[i-1][x][1]*(K-x))%=mo;
					(ssb0+=dp[i-1][x][0])%=mo;
					(ssb1+=dp[i-1][x][1])%=mo;
				}
			}
			
			FOR(j,K+1) {
				if(i==N&&j) continue;
				(dp[i+1][j][0]+=s0)%=mo;
				(dp[i+1][j][1]+=s1)%=mo;
				if(i) {
					ssb0+=mo-dp[i-1][K-j][0];
					ssb1+=mo-dp[i-1][K-j][1];
					(dp[i+1][j][0]+=ssa1)%=mo;
					(dp[i+1][j][1]+=ssa0)%=mo;
					(ssa0+=mo-ssb0)%=mo;
					(ssa1+=mo-ssb1)%=mo;
				}
			}
		}
		ll ret=dp[N+1][0][0]+mo-dp[N+1][0][1];
		cout<<ret%mo<<endl;
		
	}
}

まとめ

包除原理まではわかっても、そこからの状態遷移がちょっと難しい。