kmjp's blog

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

Codeforces #413 E. Aquarium decoration

これはややこしいけど解けてよかった。
http://codeforces.com/contest/799/problem/E

問題

N個のアイテムがあり、それぞれコストC[i]が設定されている。
これらのアイテムからM個を選択することを考える。

ここで2人の利用者がおり、それぞれ好みのアイテムのリストが与えられる。
選択したM個中、両者とも最低K個以上好みのアイテムを選ぶとき、コストの総和の最小値を求めよ。

解法

アイテム群を、両者が好みかそうでないかで4通りに分け、それぞれコストが低い順にソートする。
また、それとは別に全アイテムをコストが低い順にソートする。

両者がともに好みなアイテムを選ぶ数pを総当たりしよう。
そのとき、それぞれ片方だけが好きなものは最低(K-p)個ずつ選ばなければならない。
あとは、好き嫌い問わず残りM-(p+(K-p)*2)個をコストの小さい順に選べばよい。

事前にリストの累積和を取っておくと、「両者がともに好みなアイテムp個」「片方だけが好きなもの(K-p)個」のコストの総和は簡単に取れる。
あとは、全アイテムをコストが低い順にソートしたものについて、BITを2つ使って未取得のアイテム数およびそのコストの総和を求められるようにしておこう。
前者のBITを二分探索すれば、未取得のアイテムからM-(p+(K-p)*2)個選ぶときのコストの総和が求められる。

template<class V, int ME> class BIT {
public:
	V bit[1<<ME],val[1<<ME];
	V total(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	V add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
	V set(int e,V v) { add(e,v-val[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;
	}
};

int N,M,K;
ll C[202020];
int A[2],X[2][202020];
int mask[202020];
vector<pair<ll,int>> T[5];
vector<pair<ll,int>> P;
int id[202020];

BIT<ll,20> sum,can;
ll sum2;

void hoge(int v,int sgn) {
	sum2 += sgn*C[v];
	sum.add(id[v],-sgn*C[v]);
	can.add(id[v],-sgn*1);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	FOR(i,N) cin>>C[i];
	FOR(x,2) {
		cin>>A[x];
		FOR(i,A[x]) {
			cin>>X[x][i];
			mask[X[x][i]-1] |= 1<<x;
		}
	}
	
	FOR(i,N) {
		T[mask[i]].push_back({C[i],i});
		T[4].push_back({C[i],i});
		P.push_back({C[i],i});
	}
	FOR(i,5) sort(ALL(T[i]));
	FOR(i,N) {
		id[T[4][i].second]=i+1;
		sum.add(id[T[4][i].second],C[T[4][i].second]);
		can.add(id[T[4][i].second],1);
	}
	ll mi=1LL<<60;
	FOR(i,min(K,(int)T[3].size())) hoge(T[3][i].second,1);
	
	for(i=K;i>=0;i--) {
		if(i<K && i<T[3].size()) hoge(T[3][i].second,-1);
		j=K-i;
		if(j) {
			if(j<=T[1].size()) hoge(T[1][j-1].second,1);
			if(j<=T[2].size()) hoge(T[2][j-1].second,1);
		}
		
		if(i<=T[3].size() && j<=T[1].size() && j<=T[2].size()) {
			int left=M-(i+j+j);
			if(left>=0) {
				x=can.lower_bound(left);
				mi=min(mi,sum2+sum.total(x));
			}
		}
		
	}
	
	if(mi==1LL<<60) mi=-1;
	cout<<mi<<endl;
	
	
}

まとめ

解けてよかったけど、もう少しシンプルに書きたいな。