SRM 311 Div1 Medium / Div2 Hard FloatingMedian

データ構造のタグがないことに気づいた。

問題

長さNの数列aが与えられる。
aの連続するK個の部分数列全てについて、それぞれの中央値を合計するといくつになるか、求めよ。

ただし長さKの数列の中央値とは、(K+1)/2番目に小さい項の値を言う。

制約条件

N≦25万
K≦5000
0≦a[i]≦65535

方針

平方分割やセグメント木でk-th number的なことをやろうとするとTLE.

a[i]の値が65536通りしかないので、0以上j以下の項の数はBinary Indexed Tree使ってO(log k)で求めることができる。
よって、中央値は二分探索によりO(log 65536*log k)で求めることができる。
次の項を入れて、左端の項を出しながら、中央値を求めることをN-K+1回繰り返せばよい。

ソースコード

const int MX=300000,MX2=70000;
int a[MX],N,bit[MX2+1];
int sum(int i){
	int s=0;
	while(i>0)s+=bit[i], i-=i&-i;
	return s;
}
void add(int i,int x){
	while(i<=MX2)bit[i]+=x, i+=i&-i;
}

class FloatingMedian {
	public:
	long long sumOfMedians(int seed, int mul, int ad, int N, int K) 
	{
		::N=N;
		a[0]=seed;
		rep(i,N-1)a[i+1]=((ll)a[i]*mul+ad)%65536;
		rep(i,MX2+1)bit[i]=0;
		
		ll ans=0;
		rep(i,K-1)add(a[i]+1,1);
		rep(i,N-K+1){
			add(a[i+K-1]+1,1);
			int lo=-1,hi=65536,mid;
			while(lo+1<hi){
				mid=(lo+hi)/2;
				if(sum(mid+1)<(K+1)/2)lo=mid; else hi=mid;
			}
			ans+=hi;
			add(a[i]+1,-1);
		}
		return ans;
	}
};