CSAcademy Round #22 : E. Limited Swaps

いかにもありそうで、なかった問題。
https://csacademy.com/contest/round-22/#task/limited-swaps

問題

N要素の整数列A[i]が与えられる。
このA[i]のうち最大K回まで隣接要素をswapできるとき、得られる辞書順最大の数列を答えよ。

解法

辞書順最大のものを求める問題なので、数列の頭から確定させていこう。
その際、残されたswap可能回数の範囲で頭に持ってこれる最大要素を持って来ればよい。

あとは上記を高速に処理できるデータ構造を考える。
各要素が位置未確定かどうかを示すBITと、範囲内の最大値およびその添え字を求めるSegTreeを用いる。

最初のBITは、初期状態では0～(N-1)番目の値にそれぞれ1を格納しておく。
ある要素を答えとなる数列の頭に持っていって位置が確定すると、その要素の元の位置の値が0となる。
このBITで、和が(K+1)以下となる範囲を求めよう。その範囲が次にK回以下のswapで先頭に持ってこれる要素の範囲となる。
あとはその範囲内で、SegTreeを使い最大値を求める、ということを繰り返していけばよい。

int N;
ll K;

template<class V,int NV> class SegTree_Pair {
public:
	vector<pair<V,int> > val;
	static V const def=-1;
	pair<V,int> comp(pair<V,int> l,pair<V,int> r){
		if(r.first>l.first) return r;
		else return l;
	}
	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<ll,1<<20> st;

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	V add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
	int lower_bound(V val) {
		V tv=0; int i,ent=0;
		for(i=ME-1;i>=0;i--) if(tv+bit[ent+(1<<i)-1]<=val) tv+=bit[ent+(1<<i)-1],ent+=(1<<i);
		return ent;
	}
};
BIT<int,20> bt;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		cin>>x;
		bt.add(i,1);
		st.update(i,x);
	}
	
	FOR(i,N) {
		int R;
		
		K++;
		if(K>=bt(1<<19)) {
			R=N;
		}
		else {
			R=bt.lower_bound(K);
		}
		auto p=st.getval(0,R);
		K -= bt(p.second);
		cout<<p.first<<" ";
		bt.add(p.second,-1);
		st.update(p.second,-1);
	}
	cout<<endl;
	
	
}

まとめ

わかってしまえば簡単。