CF891C Envy

CF891C Envy
Q3.2.2.4. MST 边集询问

最小生成树(森林)中相同权值的边数量不变
先求加入第 i 条边后两个点所在连通块
对于每一个权值: 用并查集判断所在连通块是否有环(有则无解)

点击查看代码

 
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <utility>
#include <array>
#include <vector>

using namespace std;

const int N = 5e5 + 5, M = 1e6 + 5;

int n, m;
struct Edges {
	int a, b, w, id; // 端点1, 端点2, 权值, 编号
	int fa, fb; // a 所在连通块, b 所在连通块
} edges[M];
int father[N];
int pos[N]; // 编号为 i 的边的下标

int find(int x) {
	if(x == father[x]) return x;
	return father[x] = find(father[x]);
}

bool cmp(const Edges &a, const Edges &b) {
	return a.w < b.w;
}

int main() {
	scanf("%d%d", &n, &m);
	for(int i = 1, a, b, c; i <= m; i ++) {
		scanf("%d%d%d", &a, &b, &c);
		edges[i] = {a, b, c, i, 0, 0};
	}
	sort(edges + 1, edges + m + 1, cmp);
	for(int i = 1; i <= n; i ++) father[i] = i;
	for(int i = 1; i <= m; ) {
		int j = i;
		while(edges[j].w == edges[i].w) {
			edges[j].fa = find(edges[j].a); // 所在连通块
			edges[j].fb = find(edges[j].b); // 所在连通块
			j ++;
		}
		while(i < j) {
			father[find(edges[i].a)] = find(edges[i].b); // 合并
			i ++;
		}
	}
	for(int i = 1; i <= m; i ++) pos[edges[i].id] = i;
	int q;
	scanf("%d", &q);
	for(int i = 1; i <= n; i ++) father[i] = i;
	while(q --) {
		int edges_cnt;
		scanf("%d", &edges_cnt);
		vector<int> query;
		bool ok = true;
		for(int i = 1, x; i <= edges_cnt; i ++) {
			scanf("%d", &x);
			x = pos[x];
			query.push_back(x);
		}
		sort(query.begin(), query.end(), [&](const int &x, const int &y) {
			return edges[x].w < edges[y].w;
		});
		for(int i = 0, sz = query.size(); i < sz && ok; ) {
			int j = i; // [i, j)
			while(j < sz && edges[query[i]].w == edges[query[j]].w) {
				int fa = find(edges[query[j]].fa), fb = find(edges[query[j]].fb); // 环!
				if(fa == fb) {
					ok = false;
					break;
				}
				father[fa] = fb; // 合并
				j ++;
			}
			while(i < j) {
				father[edges[query[i]].fa] = edges[query[i]].fa; // 恢复
				father[edges[query[i]].fb] = edges[query[i]].fb; // 恢复
				i ++;
			}
		}
		putchar(ok ? '1' : '0');
	}
	return 0;
}