Codeforces 891 C Envy

题目大意

给定一个 \(n\) 个点 \(m\) 条边的连通的无向图，每条边有一个权值，可能有重边。给出 \(q\) 组询问，一组询问给出 \(k\) 条边，问是否存在一棵最小生成树包含这 \(k\) 条边。

思路

这道题是我独立想出来的，并没有看题解。

先考虑只有一组询问的做法。
判断某些边是否在某个最小生成树里，比较容易想到的做法是 Kruskal 算法：
为了表述方便，将输入的 \(m\) 条边构成的集合记做 \(E\)，将查询的 \(k\) 条边构成的集合记做 \(E_q\) 。
假设 \(E\) 中的边共有 \(c\) 个不同的权值，从小到大依次为 \(w_1, w_2, \dots, w_c\)。
将所有「权值小于 \(w_1\)」 的边都插入并查集，然后再将 \(E_q\) 所有权值等于 \(w_1\) 的边逐个插入并查集，看是否有矛盾（即是否有某条边的端点在同一个连通分量的情况），若有矛盾则说明不存在满足条件的最小生成树。若无矛盾则将 \(E\) 中权值在 \([w_1,w_2)\) 范围内的边插入并查集，再把 \(E_q\) 中权值为 \(w_2\) 的边插入并查集，看是否矛盾；这样依次检验下去，若始终无矛盾则说明存在满足条件的最小生成树。

再考虑如何将这个做法用到多组询问上。
不难想到可以将询问离线处理。将 \(q\) 组询问中权值相等的边放在一起处理。假设当前检验的是询问中权值为 \(w\) 的边（此时 \(E\) 中权值小于 \(w\) 边都已插入并查集），每次检验完某个询问中权值为 \(w\) 的边之后，需要将这些边从并查集中删除。

在用并查集维护（动态）图的连通性时，如果需要支持删边操作，那么可以不采用路径压缩而只采用按秩合并；这样就使得删边操作很容易实现。而按秩合并可以保证一棵 \(n\) 个点的树的高度为 \(O(\log n)\)，这样找到根节点的复杂度为 \(O(\log n)\)，而合并和删边都是 \(O(1)\) 的；总复杂度为 \(O(m\log m + K\log n)\)，其中 \(K =\sum_{1\le i\le q}k_i\) 。

用 cin cout 竟然 TLE 了。。。

#include <bits/stdc++.h>
using namespace std;

const int N = 5e5 + 5;
int par[N];
int rk[N];

int root(int x){
	while(x != par[x]) x = par[x];
	return x;
}

struct his{
	int u, rk1;
	int v, rk2;
};

void unite(int x, int y){
	if(rk[x] > rk[y])
		par[y] = x;
	else{
		par[x] = y;
		if(rk[x] == rk[y])
			++rk[y];
	}
}


void divide(const his &x){
	par[x.u] = x.u;
	par[x.v] = x.v;
	rk[x.u] = x.rk1;
	rk[x.v] = x.rk2;
}


struct edge{
	int u, v, l;
}e[N];

int b[N];


his h[N];

struct qedge{
	int eid, qid;
}a[N];

bool cmp(int x, int y){
	return e[x].l < e[y].l;
}

bool cmp2(const qedge &x, const qedge &y){
	return e[x.eid].l < e[y.eid].l || (e[x.eid].l == e[y.eid].l && x.qid < y.qid); 
}


bool ans[N];



int main(){
	int n, m;
	scanf("%d%d", &n, &m);
	for(int i = 1; i <= m; i++){
		int u, v, l;
		scanf("%d%d%d", &u, &v, &l);
		e[i] = {u, v, l};
		b[i] = i;
	}
	


	int q;
	scanf("%d", &q);
	int nqe = 0;
	for(int i = 0; i < q; i++){
		int k;
		scanf("%d", &k);
		for(int j = 0; j < k; j++){
			int x;
			scanf("%d", &x);
			a[nqe++] = {x, i};
		}
	}

	sort(a, a + nqe, cmp2);

	for(int i = 1; i <= n; i++)
		par[i] = i;

	sort(b + 1, b + m + 1, cmp);

	for(int i = 1, j = 0; i <= m && j < nqe; ){
		for(; e[b[i]].l < e[a[j].eid].l; i++){
			int r1 = root(e[b[i]].u);
			int r2 = root(e[b[i]].v);
			if(r1 != r2)
				unite(r1, r2);
		}

		for(; j < nqe && e[a[j].eid].l == e[b[i]].l; ){
			int sz = 0;
			for(int k = a[j].qid; a[j].qid == k && e[a[j].eid].l == e[b[i]].l; j++){
				if(ans[k]) continue;
				auto &x = e[a[j].eid];
				int r1 = root(x.u), r2 = root(x.v);
				if(r1 == r2)
					ans[k] = 1;
				else{
					h[sz++] = {r1, rk[r1], r2, rk[r2]};
					unite(r1, r2);
				}
			}
			for(int i = sz - 1; i >= 0; i--){
				divide(h[i]);
			}
		}
	}

	for(int i = 0; i < q; i++)
		puts(ans[i] ? "NO" : "YES");
	return 0;
}