CF 877 F - Ann and Books
問題
問題概要
長さの数列と整数が与えられる。クエリ:「閉区間内の連続部分列で、その和がになるものの個数を答える」が個与えられるのでそれぞれに答えよ。
アイデア
prefix sumを使えば、ある区間の和はで取ることが出来る。以下ではと表す。
Mo's Algorithm(詳しくは他の人の詳しいを解説を見たほうが良いが、要素の更新がない/クエリ先読み可/区間の端を1だけ伸縮させた時の値の変化を簡単に得られる 時に有効)を使って区間の左端/右端を1伸ばす/縮める時に答えがどう変化するかを更新していきながら、該当するクエリと区間が一致した時に答えを記入していく、という流れになる。
具体的に何が分かると嬉しいかを考えてみる。いま、閉区間に対する情報を持っている状況を考える。
右端を+1
右端を右に動かしたいとする。その時に、答えの増加分は(新たに増える区間は必ず右端がなので)、を満たすの個数に一致する。
式変形すると、となり、これを満たすが何個あるかは、に対してmapなどで頻度をカウントしておけばこの更新をで行える事がわかる。
そして、答えの増分を計算した後に、mapにの分を足して、完了となる。
左端を-1
左端を左に動かしたいとする。その時に、答えの増加分は(新たに増える区間は必ず左端がなので)、を満たすの個数に一致し、式変形すると、となる。
右端を-1
右端を左に動かしたいとする。その時に、答えの減少分は(消滅する区間は必ず右端がなので)、を満たすの個数に一致する。
答えの減少分を計算する前にmapからの分を除き、を満たすiの個数を数えて答えから引いておけばよい。
左端を+1
左端を右に動かしたいとする。その時に、答えの減少分は(消滅する区間は必ず左端がなので)、を満たすの個数に一致し、式変形すると、となる。よって、答えの減少分を計算する前にmapからの分を除き、減少分を計算すれば良い。
このように、似たような操作で区間を1伸ばしたり縮めたりがで出来るので、全体の計算量がとなるが、残念ながらこの問題はこれを通させてくれない。だが、このlogは頑張って座標圧縮すると消すことが出来る。
今回、クエリ全体で現れる数は、の通りしかないので、事前処理としてこれらとindexを対応させておくことで、mapの部分はただの配列に置き換えることができ、とできる。
実装(C++)
#include <bits/stdc++.h> using namespace std; using ll = long long; #define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i)) #define all(x) (x).begin(),(x).end() #define pb push_back #define fi first #define se second using pi = pair<int,int>; using P = pair<pi,int>; const int B = 320; const int Q = 100000; const int N = 100002; ll p[N]; // p[i][-k,0,k] int f[N][3]; // ((R,L),qid) vector<P> query[B]; ll ans[Q]; int L=1, R=0; int ct[3*N]={}; ll now = 0; void pr(){ ++R; now += ct[f[R][0]]; ++ct[f[R][1]]; } void mr(){ --ct[f[R][1]]; now -= ct[f[R][0]]; --R; } void pl(){ --ct[f[L-1][1]]; now -= ct[f[L-1][2]]; ++L; } void ml(){ --L; now += ct[f[L-1][2]]; ++ct[f[L-1][1]]; } int main(){ int n; ll k; cin >>n >>k; vector<ll> a(n); rep(i,n) { cin >>a[i]; if(a[i]==2) a[i] = -1; } rep(i,n){ ll t; cin >>t; a[i] *= t; } // prefix sum p[0]=0; for(int i=1; i<=n; ++i) p[i] = p[i-1]+a[i-1]; // 座標圧縮 vector<ll> kk({-k,0,k}); vector<ll> uniq_p; rep(i,n+1){ for(ll j:kk) uniq_p.pb(p[i]+j); } sort(all(uniq_p)); uniq_p.erase(unique(all(uniq_p)), uniq_p.end()); rep(i,n+1){ rep(j,3){ ll val = p[i]+kk[j]; f[i][j] = lower_bound(all(uniq_p),val) - uniq_p.begin(); } } // クエリを左端をキーにして平方分割 int q; cin >>q; rep(i,q){ int l,r; cin >>l >>r; query[l/B].pb({{r,l},i}); } ++ct[f[0][1]]; rep(i,B){ sort(all(query[i])); for(const auto &qq:query[i]){ int l = qq.fi.se, r = qq.fi.fi, qid = qq.se; while(R<r) pr(); while(r<R) mr(); while(L<l) pl(); while(l<L) ml(); ans[qid] = now; } } rep(i,q) cout << ans[i] << endl; return 0; }