kmjp's blog

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

AtCoder ARC #075 : E - Meaningful Mean

こちらは問題ないのだけどFで手間取りすぎた。
http://arc075.contest.atcoder.jp/tasks/arc075_c

問題

N要素の整数列Aと、整数Kが与えられる。
整数列の連続部分列のうち、算術平均がKを超えるものはいくつあるか。

解法

先にAの各要素からKを引いておくと、算術平均は気にしなくてよくてAの総和が非負である部分列を探せばよくなる。

累積和としてS[i]=sum(A[0...i])を計算しておこう。
連続部分列の左端をLとすると、S[R]≧S[L-1]となるR≧Lが条件を満たす右端Rとなる。

先にS[x]を求めて座標圧縮しておくと、そのようなRの数を高速に求めることができる。

int N,K;
ll A[202020];
ll S[202020];
vector<ll> V;

template<class V, int ME> class BIT {
public:
	V bit[1<<ME],val[1<<ME];
	V total(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	V add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
	V set(int e,V v) { add(e,v-val[e]);}
	int lower_bound(V val) {
		V tv=0; int i,ent=0;
		for(i=ME-1;i>=0;i--) if(tv+bit[ent+(1<<i)-1]<val) tv+=bit[ent+(1<<i)-1],ent+=(1<<i);
		return ent;
	}
};

BIT<int,20> bit;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	V.push_back(-1LL<<60);
	V.push_back(0);
	FOR(i,N) {
		cin>>A[i];
		A[i]-=K;
		S[i+1]=S[i]+A[i];
		V.push_back(S[i+1]);
	}
	sort(ALL(V));
	ll ret=0;
	int added=0;
	for(i=N-1;i>=0;i--) {
		x=lower_bound(ALL(V),S[i+1])-V.begin();
		added++;
		bit.add(x,1);
		y=lower_bound(ALL(V),S[i])-V.begin();
		ret += added-bit.total(y-1);
	}
	cout<<ret<<endl;
	
}

まとめ

ここまでは良かった。