kmjp's blog

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

Looksery Cup 2015 : F. Yura and Developers

これ系の列挙問題いつもすんなり思い浮かばないなぁ。
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でしばしば出るけどいつもすんなり解けないんだよなぁ。