kmjp's blog

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

CSAcademy Round #39 : E. K Swap

これがきっちり解けたのは良かったね。
https://csacademy.com/contest/round-39/task/k-swap/

問題

N要素の数列A[i]がある。
隣接要素同士、差の絶対値がK以下の場合swap可能とする。
swapを繰り返して得られる辞書順最小の数列を答えよ。

解法

辞書順最小を求める問題の定番テクとして、先頭から埋めていこう。
問題の条件を言い換えると、先頭から何要素か確定したとき、次の要素は以下を満たすものが来るのが良い。

  • 自信より手前に、swap不可能な未確定の値がない(swap不可能な値があるとそれを追い越せない)
  • 上記を満たすうち最小の値

これを求めるデータ構造を考えよう。
まず、各要素において自身の手前に追い越せない値が何要素あるかを数えよう。
これはBITで各値の登場回数を数えていけば求められる。

次に、数列中のindexを、対応する要素が昇順となるように並べ替える。
また、各indexに対応する手前の追い越せない値の数の最小値を求めるSegTreeを作ろう。

このSegTreeを探索すると、追い越せない数が最小(=0)でかつ最小要素をO(logN)で求めることができる。
今ここでA[x]を答えの数列における次の未確定要素として確定させたとする。
そうしたらSegTreeにおいて[1,A[x]-K-1]または[A[x]+K+1,10^5]の範囲に相当する領域をデクリメントすればよい。
(確定済みの値は、手前の追い越せない要素数をinfにしておけば2回選択されることはない)

上記処理をN回繰り返せばよい。

int N,K;
int A[101010];
int num[101010];
pair<int,int> P[101010];

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {
		if(e<0) return 0;
		if(e>200000) e=200000;
		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;}
};
BIT<int,20> bt;


template<class V,int NV> class SegTree_3 {
public:
	vector<pair<V,int> > val;
	vector<pair<V,int> > mi;
	static V const def=1010100;
	SegTree_3(){
		int i;
		val.resize(NV*2); mi.resize(NV*2);
		FOR(i,NV) val[i+NV]=mi[i+NV]=make_pair(0,i);
		FOR(i,NV) val[i]=make_pair(0,0);
		for(i=NV-1;i>=1;i--) mi[i]=min(mi[2*i],mi[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,0);
		if(x<=l && r<=y) return mi[k];
		auto a=min(getval(x,y,l,(l+r)/2,k*2),getval(x,y,(l+r)/2,r,k*2+1));
		a.first += val[k].first;
		return a;
	}
	
	void update(int x,int y, V v,int l=0,int r=NV,int k=1) {
		if(l>=r) return;
		if(x<=l && r<=y) {
			val[k].first+=v;
			mi[k].first+=v;
		}
		else if(l < y && x < r) {
			update(x,y,v,l,(l+r)/2,k*2);
			update(x,y,v,(l+r)/2,r,k*2+1);
			mi[k]=min(mi[k*2],mi[k*2+1]);
			mi[k].first += val[k].first;
		}
	}
};
SegTree_3<int,1<<20> st;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		cin>>A[i];
		num[i]=bt(200000)-bt(A[i]+K)+bt(A[i]-K-1);
		bt.add(A[i],1);
		P[i]={A[i],i};
	}
	sort(P,P+N);
	FOR(i,N) st.update(i,i+1,num[P[i].second]);
	FOR(i,N) {
		auto r=st.getval(0,N);
		x = r.second;
		
		cout<<P[x].first;
		if(i!=N-1) cout<<" ";
		else cout<<endl;
		
		st.update(x,x+1,1010101);
		y=lower_bound(P,P+N,make_pair(P[x].first+K+1,0))-P;
		st.update(y,N,-1);
		y=lower_bound(P,P+N,make_pair(P[x].first-K,0))-P;
		st.update(0,y,-1);
	}
	
	
}

まとめ

これが1発ACできたのはよかった。