kmjp's blog

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

AtCoder ABC #146 : E - Rem of Sum is Num

遅刻と中断があったので順位がそれなり。
https://atcoder.jp/contests/abc146/tasks/abc146_e

問題

N要素の整数列Aと正整数Kが与えられる。
連続するx要素の部分列で、部分列の値の総和をKで割った余りがxになるようなものは何通りか。

解法

A[i]の和をx個分取ってKで割るとx余るということは、x<Kでかる(A[i]-1)をx個とってKで割ると0になるのと同じ意味である。
よって、A[i]-1の累積和を取りつつ、mod Kが一致するものの数を数え上げていこう。

int N,K;
ll A[202020];
ll S[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	map<ll,int> M;
	ll ret=0;
	M[0]=1;
	FOR(i,N) {
		cin>>A[i];
		S[i+1]=S[i]+A[i]+K-1;
		if(i+1-K>=0) M[S[i+1-K]%K]--;
		ret+=M[S[i+1]%K];
		M[(S[i+1])%K]++;
	}
	cout<<ret<<endl;
}

まとめ

微妙にoff-by-oneエラーが怖い問題。