CF1374E2 Reading Books(hard version)

这道题是在 CF1374E1 Reading Books(easy version) 的基础上出的，而且仅仅增加了一个 \(m\) 的限制，下面的做法也是基于简单版的思路而想到的。建议先尝试简单版，然后尝试此题。

题意简述：

有 \(n\) 本书，每本书有一个阅读时间 \(t_i\)，和属性 \(a_i\)、\(b_i\)，代表是否被 Alice 和 Bob 喜欢。要求两人恰好读 \(m\) 本书，其中必须有至少 \(k\) 本被 Alice 喜欢，至少 \(k\) 本被 Bob 喜欢，求最小的总阅读时间，并输出读书方案。总阅读时间为这 \(m\) 本书的时间和。

思路：

简单版

先不考虑 \(m\)。我们可以把书分为四类：被 Alice 喜欢，被 Bob 喜欢，两人都喜欢，两人都不喜欢。我们设这四类分别为 \(a\)、\(b\)、\(ab\)、\(none\)。显然，两人都不喜欢的书没必要去选，满足了 \(k\) 的限制后，也没必要去多选书。因此，问题转换为了要选多少 \(ab\) 与 \(a\)、\(b\)。每多选一本 \(ab\) 就可以少选一本 \(a\) 和和一本 \(b\)，所以我们可以枚举选 \(i\) 本 \(ab\)，那么 \(a\) 和 \(b\) 的数量就是 \(k-i\)。我们只需要知道阅读时间前 \(i\) 小的 \(ab\) 的阅读时间和以及前 \(k-i\) 小的 \(a\) 和 \(b\) 的阅读时间和即可。这三类书的选择相互独立，因此可以分别排序，求前缀和即可。

困难版

和简单版唯一的不同在于，简单版可以读任意数量的书，而困难版必须读 \(m\) 本。也就是说，如果已经满足了 \(k\) 的限制后，读的书数量不够多，必须从未选择的书中再选出几本。

我们来分析一下刚才过程中，未选择的书有哪些。还是设选择 \(i\) 本 \(ab\)，设 \(ab\)、\(a\)、\(b\)、\(none\) 四种书的总数为 \(totab\)、\(tota\)、\(totb\)、\(totn\)。

• \(ab\) 中有 \(totab - i\) 本未被选择；
• \(a\) 中有 \(tota - k + i\) 本书未被选择；
• \(b\) 中有 \(totb - k + i\) 本书未被选择；
• \(none\) 中的我们都没有选。

我们设没有选的书的集合为 \(S\)，那么我们只需要找到 \(S\) 中阅读时间前 \(m-k*2+i\) 小的书的阅读时间和。这个可以用平衡树做。

我们来理一下这个过程。首先从小到大枚举选 \(i\) 本 \(ab\)，然后维护 \(S\) 内的元素，计算答案，如果答案可以更新，则记录一下得到答案时 \(i\) 是多少（因为还要输出方案）。

至于输出方案，首先 \(ab\)、\(a\)、\(b\) 都好处理，从小到大输出即可；对于 \(S\) 中我们选择的元素，可以先恢复 \(S\) 在得到答案的 \(i\) 时的状态，然后把得到答案的那部分输出出来即可。

总时间复杂度 \(O(n \log n)\)。细节见代码注释。

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 2e5+100;

inline int read(){
	int x = 0; char ch = getchar();
	while(ch<'0' || ch>'9') ch = getchar();
	while(ch>='0'&&ch<='9') x = x*10+ch-48, ch = getchar();
	return x;
}

struct xwx {
	ll val;
	int id;int v;
	bool operator < (const xwx &b ) const{
		return val < b.val;
	}
};// 用来存储 a、b、ab、none 中的书
struct qwq{
	int v, id;
	bool operator <(const qwq &b) const{
		if(v == b.v){
			return id < b.id;
		}
		return v < b.v;
	}
	bool operator >(const qwq &b) const{
		if(v == b.v){
			return id > b.id;
		}
		return v >b.v;
	}
};//作为 FHQ 上元素的权值，重载运算符便于删除操作。
xwx a[N], b[N], ab[N], no[N];//A喜欢，B喜欢，都喜欢，都不喜欢。

struct node{
	qwq val;
	ll sum;//sum 记录子树内的权值和。
	int siz;
	int ls, rs;
	int rnd;
};
int root;
mt19937 getrand(time(0));
struct FHQ_Treap{
	node tree[N];
	int idx;
	#define ls(x) tree[x].ls
	#define rs(x) tree[x].rs
	void push_up(int tr){
		tree[tr].sum = tree[ls(tr)].sum+tree[rs(tr)].sum+tree[tr].val.v;
		tree[tr].siz = tree[ls(tr)].siz+tree[rs(tr)].siz+1;
	}
	int New(qwq val){
		tree[++idx] = (node){val, val.v, 1, 0, 0, (int)getrand()};
		return idx;
	}
	void split_by_size(int pos, int &l, int &r, int K){//按大小分裂，也就是取出前 K 小的元素
		if(!pos) return l = r = 0, void();
		int tmp = tree[ls(pos)].siz+1;
		if(K>=tmp){
			l = pos;
			split_by_size(rs(pos), rs(pos), r, K-tmp);
			push_up(l);
		} else{
			r = pos;
			split_by_size(ls(pos), l, ls(pos), K);
			push_up(r);
		}
	}
	void split_by_val(int pos, int &l, int &r, qwq K){//按权值大小分裂，主要用于插入元素。
		if(!pos) return l = r = 0, void();
		if(K>tree[pos].val){
			l = pos;
			split_by_val(rs(pos), rs(pos), r, K);
			push_up(l);
		} else{
			r = pos;
			split_by_val(ls(pos), l, ls(pos), K);
			push_up(r);
		}
	}
	int merge(int l, int r) {
		if((!l) || (!r)) return l | r;
		if(tree[l].rnd < tree[r].rnd){
			rs(l) = merge(rs(l), r);
			push_up(l);
			return l;
		} else{
			ls(r) = merge(l, ls(r));
			push_up(r);
			return r;
		}
	}
	void insert(int val, int id){
		int dl, dr;
		split_by_val(root, dl, dr, (qwq){val, id});
		root = merge(merge(dl, New((qwq){val, id})), dr);
	}
	void del(int val, int id){
		int dl, md, dr;
		split_by_val(root, dl, dr, (qwq){val, id});
		split_by_val(dr, md, dr, (qwq){val, id+1});//将目标节点分裂出来，合并两边。
		root = merge(dl, dr);
	}
	ll query(int K){
		int dl, dr;
		split_by_size(root, dl, dr, K);//分裂出前 K 小的元素。
		ll tmp = tree[dl].sum;
		root = merge(dl,dr);
		return tmp;
	}
	void clear(){
		root = 0;
		idx = 0;
		memset(tree, 0, sizeof(tree));
	}
	void output(int pos){
		if(ls(pos)){
			output(ls(pos));
		}
		printf("%d ", tree[pos].val.id);
		if(rs(pos)){
			output(rs(pos));
		}
	}//遍历子树，输出方案
	void print(int K){
		int dl,dr;
		split_by_size(root, dl, dr, K);
		if(dl) output(dl);
	}
}s;
int ta, tb, tab, tn;
int n, m, K;
int main(){
	n = read(), m = read(), K = read();
	for(int i = 1, ar, br, t; i<=n; ++i){
		t = read(), ar = read(), br = read();
		if(ar && br){
			ab[++tab] = (xwx){t, i, t};
			s.insert(t, i);
		} else if(ar){
			a[++ta] = (xwx){t, i, t};
		} else if(br){
			b[++tb] = (xwx){t, i, t};
		} else{
			no[++tn] = (xwx){t, i, t};
			s.insert(t, i);//没人喜欢的直接插入到 S 里
		}
	}
	sort(a+1, a+ta+1);
	sort(b+1, b+tb+1);
	sort(ab+1, ab+tab+1);
	// sort(no+1, no+tn+1);
	for(int i = 1; i<=ta; ++i){
		a[i].val+=a[i-1].val;
	}
	for(int i = 1; i<=tb; ++i){
		b[i].val+=b[i-1].val;
	}
	for(int i = 1; i<=tab; ++i){
		ab[i].val+=ab[i-1].val;
	}
	// 求前缀和。
	ll ans = 0x3f3f3f3f3f3f3f3f;
	int num = -1;//记录答案在哪里产生
	int tota = ta, totb = tb, totab = tab;
	for(int i = 0; i<=min(K, tab); ++i){
		if(i > 0){
			s.del(ab[i].v, ab[i].id);
		}//枚举到的 ab 必须要选，所以就从 S 中删除。
		if((K-i)>ta || (K-i)>tb || K*2-i > m) continue;

		while(ta > (K-i)){
			s.insert(a[ta].v, a[ta].id);
			--ta;
		}
		while(tb > (K-i)){
			s.insert(b[tb].v, b[tb].id);
			--tb;
		}//将选不上的 a 和 b 都加入 S

		ll tmp = ab[i].val + a[ta].val + b[tb].val + s.query(m-K*2+i);
		if(ans > tmp){
			num = i;
			ans = tmp;
		}
	}
	if(ans == 0x3f3f3f3f3f3f3f3f){
		puts("-1");
		return 0;
	}
	printf("%lld\n", ans);
	s.clear();
	for(int i = num+1; i<=totab; ++i){
		s.insert(ab[i].v, ab[i].id);
	}
	for(int i = (K-num+1); i<=tota; ++i){
		s.insert(a[i].v, a[i].id);
	}
	for(int i = (K-num+1); i<=totb; ++i){
		s.insert(b[i].v, b[i].id);
	}
	for(int i = 1; i<=tn; ++i){
		s.insert(no[i].v, no[i].id);
	}//将 S 恢复到 i = num 时的状态

	for(int i = 1; i<=num; ++i){
		printf("%d ", ab[i].id);
	}
	for(int i = 1; i<=K-num; ++i){
		printf("%d ", a[i].id);
	}
	for(int i = 1; i<=K-num; ++i){
		printf("%d ", b[i].id);
	}//顺序输出 ab、a、b
	s.print(m-K*2+num);//输出 S 中选择的元素。
	return 0;
}