kmjp's blog

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

Donutsプロコンチャレンジ2015 : D - ドーナツの箱詰め

ちょっと手こずったけどなんとか自力で完答。
http://donuts-2015.contest.atcoder.jp/tasks/donuts_2015_4

問題

N個の箱があり、それぞれの箱の体積はC[i]である。
一部の箱を入れ子にすることで、全体をK個の箱にしたい。
箱を入れ子にする際、その体積の差分だけの緩衝材が必要となる。

ここでQ個のクエリがある。
各クエリでは箱の番号が指定されており、その箱は以後使用しないことにする。
各クエリ毎に、K個の箱を作るのに必要な緩衝材の最小値を求めよ。

解法

まず緩衝材の最小値を考える。
C[i]を体積順にソートしたものとする。
K==1の場合、必要な緩衝材は(C[1]-C[0])、(C[2]-C[1])、...、(C[N-1]-C[N-2])の総和である。
K==2の場合、C[0]~C[x]を1個の箱、C[x+1]~C[N-1]をもう1個の箱とすると(C[x+1]-C[x])の部分の緩衝材を不要とすることができる。
一般にKに対し、(C[1]-C[0])、(C[2]-C[1])、...、(C[N-1]-C[N-2])のうち(K-1)箇所の緩衝材を不要とすることができる。

だとすると、当然ながらこれらの値のうち大きさ上位(K-1)個の緩衝材を不要としたい。
よってクエリ毎にこの上位(K-1)個を更新していけば良い。

これには上位(K-1)個の緩衝材の体積を保持するmultisetと、それ以外の緩衝材の体積を保持するmultiset、および後者のmultisetの総和を管理するとよい。

箱を1つ取り除く処理では、(体積ソート後で)i番の箱より1つ小さな箱の番号をL[i]、1つ大きな箱の番号をR[i]を双方向リンクリストのように管理する。
i番の箱を取り除く場合は緩衝材の体積のmultisetから(C[R[i]]-C[i])と(C[i]-C[L[i]])を取り除いて
(C[R[i]]-C[L[i]])を加え、さらにリンクリストを更新していけばよい。

int N,K;
int C[202000];
pair<int,int> P[202000];
int Q,D[202000];
int L[202000],R[202000];
ll tot=0;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		cin>>C[i];
		P[i]=make_pair(C[i],i);
	}
	sort(P,P+N);
	multiset<int> sm,la;
	FOR(i,N) {
		L[P[i].second]=(i==0)?-1:P[i-1].second;
		R[P[i].second]=(i==N-1)?N:P[i+1].second;
		if(i!=N-1) la.insert(P[i+1].first-P[i].first);
	}
	
	while(la.size()>=K) tot+=*la.begin(),sm.insert(*la.begin()), la.erase(la.begin());
	cout<<tot<<endl;
	cin>>Q;
	FOR(i,Q) {
		multiset<int>::iterator mit;
		cin>>D[i],D[i]--;
		if(L[D[i]]==-1) {
			int dif=C[R[D[i]]]-C[D[i]];
			mit=la.find(dif);
			if(mit!=la.end()) la.erase(mit);
			else tot-=dif, sm.erase(sm.lower_bound(dif));
			L[R[D[i]]]=-1;
		}
		else if(R[D[i]]==N) {
			int dif=C[D[i]]-C[L[D[i]]];
			mit=la.find(dif);
			if(mit!=la.end()) la.erase(mit);
			else tot-=dif, sm.erase(sm.lower_bound(dif));
			R[L[D[i]]]=N;
		}
		else {
			int dif1=C[R[D[i]]]-C[D[i]];
			int dif2=C[D[i]]-C[L[D[i]]];
			int dif3=C[R[D[i]]]-C[L[D[i]]];
			mit=la.find(dif1);
			if(mit!=la.end()) la.erase(mit);
			else tot-=dif1, sm.erase(sm.lower_bound(dif1));
			mit=la.find(dif2);
			if(mit!=la.end()) la.erase(mit);
			else tot-=dif2, sm.erase(sm.lower_bound(dif2));
			
			la.insert(dif3);
			while(la.size()>=K) tot+=*la.begin(),sm.insert(*la.begin()), la.erase(la.begin());
			L[R[D[i]]]=L[D[i]];
			R[L[D[i]]]=R[D[i]];
		}
		while(la.size()<K-1) {
			multiset<int>::iterator mit=sm.end();
			mit--;
			tot-=*mit,la.insert(*mit), sm.erase(mit);
		}
		cout<<tot<<endl;
	}
}

まとめ

multisetのfindやeraseが若干重いので連発しすぎるとTLEするので注意。