Codeforces 1166F 并查集 启发式合并

题意：给你一张无向图，无向图中每条边有颜色。有两种操作，一种是询问从x到y是否有双彩虹路，一种是在x到y之间添加一条颜色为z的边。双彩虹路是指：如果给这条路径的点编号，那么第i个点和第i - 1个点相连的边与第i个点和第i + 1个点相连的边颜色一样，其中i是偶数。

思路：这个问题相当于初了最后一步没有限制以外（最后一步只走一条边），每一步都要走颜色相同的两条边。我们先不考虑最后一步的问题，只考虑每一步都要走颜色相同的两条边，那么我们只需要扫描每一个点的出边，如果有两条边颜色相同，就在一张新图上把这两条边除这个点以外的端点相连，这样询问就变成了判断两个点是否在一个连通块中。因为只是判断是否在同一连通块中，我们不需要建图，用并查集就可以了。现在需要考虑最后一步的问题，我们可以用一个set维护与这个节点直接相邻的点，这样对于询问x, y，如果y在连通块x的相邻节点中，也可以。对于加边的操作，相当于并查集的合并，但是同时需要合并的还有set, 我们合并的时候需要判断一下两个集合的大小，把小的插入大的里面，这样可以保证每次集合合并复杂度是O(logn * logn)的。

代码：

#include <bits/stdc++.h>
using namespace std;
const int maxn = 100010;
set<int> s[maxn];
set<int>::iterator it1;
set<pair<int, int> > edge[maxn];
set<pair<int, int> >::iterator it;
int f[maxn];
int get(int x) {
	if(x == f[x]) return x;
	return f[x] = get(f[x]);
}
void merge(int x, int y) {
	int x1 = get(x), y1 = get(y);
	if(x1 == y1) return;
	if(s[x1].size() < s[y1].size()) swap(x1, y1);
	for (it1 = s[y1].begin(); it1 != s[y1].end(); it1++) {
		s[x1].insert(*it1);
	}
	s[y1].clear();
	f[y1] = x1;
}
void add(int x, int y, int z) {
	s[get(x)].insert(y);
	it = edge[x].lower_bound(make_pair(z, -1));
	if(it == edge[x].end() || it -> first != z) {
		edge[x].insert(make_pair(z, y));
	} else {
		merge(it -> second, y);
	}
}
int main() {
	int x, y, z, n, m, c, T;
	char op[5];
	scanf("%d%d%d%d", &n, &m, &c, &T);
	for (int i = 1; i <= n; i++) f[i] = i;
	for (int i = 1; i <= m; i++) {
		scanf("%d%d%d", &x, &y, &z);
		add(x, y, z);
		add(y, x, z);
	}
	while(T--) {
		scanf("%s", op + 1);
		if(op[1] == '+') {
			scanf("%d%d%d", &x, &y, &z);
			add(x, y, z);
			add(y, x, z);
		} else {
			scanf("%d%d", &x, &y);
			if(get(x) == get(y)) {
				printf("Yes
");
			} else if(s[get(x)].find(y) != s[get(x)].end()) {
				printf("Yes
");
			} else {
				printf("No
");
			}
		}
	}
}