Codeforces #442 Div2 F. Ann and Books

これ想定解でいいのかな。
http://codeforces.com/contest/877/problem/F

問題

N個の本がある。
各本には整数A[i]と種類T[i]が与えられる。
種類T[i]は1(数学)か2(経済)のいずれかである。

クエリ[L,R]が与えられるので、下記に答えよ。

• 本のうち、[L...R]に含まれる区間のうち、数学の本のA[i]の総和から経済の本のA[i]の総和を引いたものがKになるものは何通りか。

解法

まず本の種類が2種類あるのでそれを処理しよう。
これは経済の本についてA[i]を負にしておけば、結局区間内のA[i]の総和がKとなるものを求める問題となり簡単になる。
また、総和を求める問題なので累積和S[i]=S[i-1]+A[i]を求めておこう。

あとはMo's Algorithmで解いていく。
まずS[i]を、map等使い同じ値同士分類しよう。keyがS[i]、valueがiの昇順vectorとなるようにしておく。
あとは[L...R]に対する解の値から、前後1つ伸ばすか減らす場合の解の増減分を求めればよい。
例えば[L...R]から[L...(R+1)]になるとき、S[R+1]-S[x]=kとなるようなk∈[L..R]の数だけ解は増加する。
これはkeyを(S[R+1]-k)とするvalueのvectorのうち、[L...R]の範囲のものを求めればよく、二分探索で求められる。

まとめると、O(Q√(N+Q)logN)程度で求められる。
若干定数倍が厳しいので、同じ値の計算を繰り返さないよう前処理しておくとよい。

int N,K,Q;
int A[101010];
ll S[101010];
map<ll,vector<int>> V;
int L[101010],R[101010];
const int DI=350;
ll ret[101010];
vector<pair<int,int>> ev[404];

vector<int>* VR[101010];
vector<int>* VL[101010];
vector<int>::iterator VLp[101010];
vector<int>::iterator VRp[101010];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	scanf("%d%d",&N,&K);
	for(i=1;i<=N;i++) scanf("%d",&A[i]);
	V[0].push_back(0);
	for(i=1;i<=N;i++) {
		scanf("%d",&x);
		if(A[i]==1) A[i]=x;
		else A[i]=-x;
		S[i]=S[i-1]+A[i];
		V[S[i]].push_back(i);
	}
	FOR(i,N+1) {
		if(V.count(S[i]-K)) {
			VR[i]=&V[S[i]-K];
			VRp[i]=lower_bound(VR[i]->begin(),VR[i]->end(),i);
		}
		if(V.count(S[i]+K)) {
			VL[i]=&V[S[i]+K];
			VLp[i]=lower_bound(VL[i]->begin(),VL[i]->end(),i+1);
		}
		
	}
	
	scanf("%d",&Q);
	FOR(i,Q) {
		scanf("%d%d",&L[i],&R[i]);
		ev[L[i]/DI].push_back({R[i],i});
	}
	
	FOR(i,350) if(ev[i].size()) {
		sort(ALL(ev[i]));
		
		int LL=i*DI,RR=LL;
		ll num=0;
		FORR(e,ev[i]) {
			int LE=L[e.second]-1,RE=e.first;
			while(RR<RE) {
				RR++;
				if(VR[RR]) num += VRp[RR]-lower_bound(VR[RR]->begin(),VR[RR]->end(),LL);
			}
			while(LL<LE) {
				if(VL[LL]) num -= lower_bound(VL[LL]->begin(),VL[LL]->end(),RR+1)-VLp[LL];
				LL++;
			}
			while(LE<LL) {
				LL--;
				if(VL[LL]) num += lower_bound(VL[LL]->begin(),VL[LL]->end(),RR+1)-VLp[LL];
			}
			ret[e.second]=num;
		}
	}
	
	
	
	FOR(i,Q) cout<<ret[i]<<endl;
}

まとめ

解法は割とすんなり思いついたけど、この計算量で通るとは思わなかった。