これ系の列挙問題いつもすんなり思い浮かばないなぁ。
http://codeforces.com/contest/549/problem/F
問題
N要素の整数列A[i]と整数Kが与えられる。
A[i]のうち、以下の条件を満たす連続部分列はいくつあるか。
- 要素数が2以上である。
- 部分列中、最大要素の値を1つ取り除くと、残りの要素の和はKの倍数である。
解法
部分列A[L..R]内に含まれる条件を満たす連続部分列を再帰的に列挙していく。
もちろん最終的にA[1..N]に対する値が求められればそれが解となる。
事前にSegTreeを用いてAのRange Maximum Queryを求められるようにしておく。
これを用いてA[L..R]の最大値を求め、それをA[M]とする。
(同じ値が複数ある場合、どれか1つ選ぶ)
この時、A[L..R]に含まれる条件を満たす部分列の数は下記の和である。
- A[L..R]にA[M]を含み、条件を満たす部分列の数
- A[L..(M-1)]含まれる、条件を満たす部分列の数
- A[(M+1)..R]含まれる、条件を満たす部分列の数
下2つは単に再帰的に求めればよいので、結局一番上を求める手段があればよい。
求めたいのは、L≦i<M<j≦Rとなる(i,j)のうち(sum(A[i..j])-A[M])%K==0となるものである。
部分列の和は事前に累積和を計算しておけば簡単に求められる。
ただ、(i,j)を総当たりすると最悪ケースでO(N^2)かかるので何とかしたい。
そこで、iとjどちらか少ない方だけ全探索し、もう片方を高速に数え上げる方法を考える。
以下、iの方が少ない場合(M-L<R-M)の場合を考える。jの方が小さい場合も同様に考えられる。
S[i]=sum(A[1..i])とする。
iが固定されている場合、(sum(A[i..j])-A[M])%K==0より(S[j]-S[i-1]-A[M])%K=0なのでS[i-1]+A[M]=S[j] (mod K)となるjで、かつM<j≦Rである物を数え上げたい。
よって、事前にS[j]をS[j]%Kの値で分類し、jを配列に突っ込んでおこう。
あとは(S[i-1]+A[M])%K番の配列にあるM<j≦Rとなるjの値はlower_boundなどでlogNで求められる。
総当たりの際はi==Mのコーナーケースに注意。要素数が2以上でなければならないので、i==Mならj≧M+1でなければならない。
template<class V,int NV> class SegTree_Pair { public: vector<pair<V,int> > val; static V const def=0; pair<V,int> comp(pair<V,int> l,pair<V,int> r){ return max(l,r);} SegTree_Pair(){ val.resize(NV*2); int i; FOR(i,NV) val[i+NV]=make_pair(def,i); for(i=NV-1;i>=1;i--) val[i]=comp(val[2*i],val[2*i+1]); }; pair<V,int> getval(int x,int y,int l=0,int r=NV,int k=1) { if(r<=x || y<=l) return make_pair(def,NV); if(x<=l && r<=y) return val[k]; return comp(getval(x,y,l,(l+r)/2,k*2),getval(x,y,(l+r)/2,r,k*2+1)); } void update(int entry, V v) { entry += NV; val[entry]=make_pair(v,entry-NV); while(entry>1) entry>>=1, val[entry]=comp(val[entry*2],val[entry*2+1]); } }; SegTree_Pair<int,1<<20> st; int N,K; int A[303030]; ll S[303030]; vector<int> P[1010000]; ll dodo(int L,int R) { if(R-L<=0) return 0; auto ma=st.getval(L,R+1); ll ret=0; int i; if(ma.second-L<R-ma.second) { for(i=L;i<=ma.second;i++) { int v=(S[i-1]+ma.first)%K; ret += lower_bound(ALL(P[v]),R+1)-lower_bound(ALL(P[v]),ma.second+(i==ma.second)); } } else { for(i=ma.second;i<=R;i++) { int v=(S[i]-ma.first)%K; ret += lower_bound(ALL(P[v]),ma.second-(i==ma.second))-lower_bound(ALL(P[v]),L-1); } } return ret+dodo(L,ma.second-1)+dodo(ma.second+1,R); } void solve() { int i,j,k,l,r,x,y; string s; cin>>N>>K; P[0].push_back(0); FOR(i,N) { cin>>A[i+1]; S[i+1]=S[i]+A[i+1]; P[S[i+1]%K].push_back(i+1); st.update(i+1,A[i+1]); } cout<<dodo(1,N)<<endl; }
まとめ
こういう分割統治法、CFでしばしば出るけどいつもすんなり解けないんだよなぁ。