kmjp's blog

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

yukicoder : No.2215 Slide Subset Sum

思いつけて良かった。
https://yukicoder.me/problems/no/2215

問題

N要素の整数列Aが与えられる。
連続するM要素それぞれにおいて、Subsetのうち和がKの倍数となるものは何通りか答えよ。

解法

あいにく戻すDPでは対処できない。
Aを、M要素ごとに区切り、それぞれ前からと後ろからDPでKで割った余りごとにSubsetの組み合わせを数えよう。

各連続M要素においては、上記結果を1つまたは2つ組み合わせることで、それぞれO(K)答えられる。

int N,M,K;
int A[402020];

int L[404040][101];
int R[404040][101];

const ll mo=998244353;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	FOR(i,N) cin>>A[i];
	
	if(M==1) {
		FOR(i,N) {
			cout<<(A[i]==0)<<endl;
		}
		return;
	}
	
	int ON=N;
	N=(N+M-1)/M*M;
	FOR(i,N) {
		if(i%M==0) {
			ZERO(L[i]);
			L[i][0]=1;
		}
		FOR(j,K) {
			(L[i+1][j]+=L[i][j])%=mo;;
			(L[i+1][(j+A[i])%K]+=L[i][j])%=mo;
		}
	}
	for(i=N;i>=1;i--) {
		if(i%M==0) {
			ZERO(R[i]);
			R[i][0]=1;
		}
		FOR(j,K) {
			(R[i-1][j]+=R[i][j])%=mo;;
			(R[i-1][(j+A[i-1])%K]+=R[i][j])%=mo;
		}
	}
	for(i=0;i+M<=ON;i++) {
		ll ret=mo-1;
		if(i%M==0) {
			FOR(j,K) ret+=1LL*L[i+1][j]*R[i+1][(K-j)%K]%mo;
		}
		else {
			FOR(j,K) ret+=1LL*R[i][j]*L[i+M][(K-j)%K]%mo;
		}
		cout<<ret%mo<<endl;
	}
	
}

まとめ

しばらく悩んでしまった。