CSAcademy Round #10 : D. Subarray Medians

一見苦戦するかと思ったけど意外に解けた。
https://csacademy.com/contest/round-10/#task/subarray-medians

問題

N要素の数列Aが与えられる。
Aのうち連続する奇数長の部分列A[i...j]を考える。
i*j*(A[i..j]のmedian)の総和を求めよ。

解法

各A[x]がmedianになるケースを考える。

A[j](j＞x)を辿り、A[(x+1)...j]の中にあるA[x]以上の要素数と、未満の要素数の差を順次求め、それをd(x,j)としよう。
S(x,D) := d(x,j)==Dとなるjの総和 とする。

次に同様にA[i](i＜x)を辿り、A[i...(x-1)]の中にあるA[x]より大きい要素数と、以下の要素数差を順次求め、それをd(x,i)としよう。
d(x,i) = -d(x,j)であるようなA[i...j]はA[x]をmedianとし条件を満たすので、A[x]*i*S(x,-d(x,i))の総和を求めればそれが解。

int N;
ll A[10101];
ll hoge[10101];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	
	ll ret=0;
	FOR(i,N) {
		ZERO(hoge);
		y=5000;
		for(x=i;x<N;x++) {
			if(A[x]>A[i]) y++;
			else y--;
			if(y>=0 && y<=10000) hoge[y] += x+1;
		}
		y=5000;
		for(x=i;x>=0;x--) {
			if(A[x]>=A[i]) y--;
			else y++;
			if(y>=0 && y<=10000) ret += A[i]*(1+x)*hoge[y];
		}
	}
	cout<<ret<<endl;
}

まとめ

かなり迷ったけど解けて良かった。