yukicoder : No.2956 Substitute with Average

コードは意外に短い。
https://yukicoder.me/problems/no/2956

問題

整数列Aが与えられる。
Aのうち連続部分列を指定し、その中を平均値で置き換えることを考える。
こうして得られるAは何通りか。

解法

区間のうち両端の値が更新されないケースを数え上げることにする。
A[L...R]を平均値で置き換えるとき、A[L]とA[R]のいずれかが平均値と一致すれば良い。
B[i]=A[i]-xとすると、A[L...R]の平均値がxであれば、B[L..R]の総和が0で、B[L]またはB[R]=0となるケースを数え上げよう。

Aは1～30しか含まれないので、x=1～30を総当たりする。
右端Rを走査しながら、sum(B[0...R])=sum(B[0...(L-1)])となるLのうち、B[R]=0ならB[L]を問わず、B[R]≠0ならB[L]=0となるものを数え上げればよい。

int N;
int A[303030],B[202020],S[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	ll ret=1LL*N*(N+1)/2+1;
	
	for(x=1;x<=30;x++) {
		unordered_map<int,int> M;
		unordered_map<int,int> M0;
		
		M[0]++;
		FOR(i,N) {
			B[i]=A[i]-x;
			S[i+1]=S[i]+B[i];
			
			if(B[i]==0) {
				ret-=M[S[i+1]];
				M0[S[i]]++;
			}
			else {
				ret-=M0[S[i+1]];
			}
			M[S[i+1]]++;
		}
	}
	
	cout<<ret<<endl;
}

まとめ

一見ややこしいけどコードが意外に短いのはいいな。