yukicoder : No.919 You Are A Project Manager

意外と実装が面倒な問題だった。
https://yukicoder.me/problems/no/919

問題

N要素の整数列が与えられる。

以下の処理を行うことを考える。
まず任意の正整数Kを決める。
数列の長さがK以上の場合、先頭K個か末尾K個を抜き出し、抜き出した要素の中央値の総和を求める。
まだ数列の長さがK以上残っていても、途中でやめてもよい。

中央値の総和の最大値を求めよ。

解法

Kを総当たりする。
仮に、先頭または末尾からK個ずつ抜き出した要素の中央値がすでに分かっているとする。
L(n,K) := 数列中先頭からK個の要素をn回抜き出した場合の中央値の総和
R(n,K) := 数列中末尾からK個の要素をn回抜き出した場合の中央値の総和

このKに対しては、max(L(n,K)+R(m,K)) (n+m≦N/K)が解の候補となる。この値の列挙はO(N/K)で済むので、Kを総当たりしてもO(N*f(N)) (f(N)は調和級数)で済む。

あとは、各区間における中央値を求めることを考える。愚直に行うとO(N^2)かかるが、Mo's Algorithmを使うとO(N√NlogN)で済む。
先端・末尾が伸縮する数列における中央値は、値の下半分と上半分を格納する2つのmultisetを使うと容易に求められる。

int N;
int A[10101];
map<int,int> T[101010];
vector<pair<int,int>> E[101];

multiset<int> X,Y;
void add(int v) {
	X.insert(v);
	Y.insert(*X.rbegin());
	X.erase(X.find(*X.rbegin()));
	X.insert(*Y.begin());
	Y.erase(Y.begin());
	if(X.size()>Y.size()+1) {
		Y.insert(*X.rbegin());
		X.erase(X.find(*X.rbegin()));
	}
}

void del(int v) {
	if(X.count(v)) {
		X.erase(X.find(v));
		if(X.size()<Y.size()) {
			X.insert(*Y.begin());
			Y.erase(Y.begin());
		}
	}
	else {
		Y.erase(Y.find(v));
		if(X.size()>Y.size()+1) {
			Y.insert(*X.rbegin());
			X.erase(X.find(*X.rbegin()));
		}
	}
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	for(l=1;l<=N;l++) {
		for(x=0;x+l<=N;x+=l) E[x/100].push_back({x+l,x});
		for(x=N-l;x>=0;x-=l) E[x/100].push_back({x+l,x});
	}
	FOR(i,100) {
		sort(ALL(E[i]));
		int L=i*100,R=i*100;
		X.clear(),Y.clear();
		
		FORR(e,E[i]) {
			while(e.second<L) add(A[--L]);
			while(R<e.first) add(A[R++]);
			while(L<e.second) del(A[L++]);
			while(e.first<R) del(A[--R]);
			
			assert(X.size()==Y.size() || X.size()==Y.size()+1);
			T[e.second][e.first]=*X.rbegin();
		}
	}
	
	ll ret=0;
	for(l=1;l<=N;l++) {
		vector<ll> R;
		R.push_back(0);
		ll sum=0,ma=0;
		for(x=N-l;x>=0;x-=l) {
			sum+=T[x][x+l];
			ma=max(ma,sum);
			R.push_back(ma);
		}
		reverse(ALL(R));
		ret=max(ret,l*R[0]);
		sum=0;
		for(x=0;x+l<=N;x+=l) {
			sum+=T[x][x+l];
			ret=max(ret,l*(sum+R[x/l+1]));
		}
	}
	cout<<ret<<endl;
}

まとめ

Mo's Algorithm使うこと自分で思いつけるかなぁ…？