Codeforces 506D Mr. Kitayuta's Colorful Graph（分块 + 并查集）

题目链接  Mr. Kitayuta's Colorful Graph

把每种颜色分开来考虑。

所有的颜色分为两种：涉及的点的个数 $> \sqrt{n}$    涉及的点的个数 $<= \sqrt{n}$

对于第一种颜色，并查集缩点之后对每个询问依次处理过来若两点连通则答案加一。

对于第二种颜色，并查集缩点之后对该颜色涉及的所有点两两之间判断是否连通，

若连通则另外开一个$map$记录答案。

最后把两个部分的答案加起来即可。

细节问题  由于每种颜色处理完之后并查集都要重新初始化，对于第一种颜色的做法，只要$memset$即可。

第二种颜色总数可能较多，所以把之前并查集的操作撤销即可。

#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b)	for (int i(a); i <= (b); ++i)
#define dec(i, a, b)	for (int i(a); i >= (b); --i)
#define MP		make_pair
#define fi		first
#define se		second

typedef pair <int, int> PII;

const int N = 1e5 + 10;

struct node{ int x, y, ans; } q[N];

unordered_map <int, int> mp[N];
map <PII, int> id, ret;
vector <int> v[N];
vector <node> e[N];
stack  <PII> s;
int n, m, qu, cnt, line;
int father[N], c[N];

int getfather(int x){ return father[x] ? father[x] = getfather(father[x]) : x; }

int gf(int x){
	if (father[x]){
		s.push(MP(x, father[x]));
		father[x] = getfather(father[x]);
		return father[x];
	}

	else return x;
}

int main(){

	scanf("%d%d", &n, &m);
	rep(i, 1, m){
		int x, y, z;
		scanf("%d%d%d", &x, &y, &z);
		if (x > y) swap(x, y);
		mp[z][x] = mp[z][y] = 1;
		e[z].push_back({x, y});
	}

	scanf("%d", &qu);
	rep(i, 1, qu){
		scanf("%d%d", &q[i].x, &q[i].y);
		if (q[i].x > q[i].y) swap(q[i].x, q[i].y);
		id[MP(q[i].x, q[i].y)] = i;
	}

	rep(i, 1, m){
		int now = (int)mp[i].size();
		if (now > 0) v[now].push_back(i);
	}

	line = sqrt(n);
	rep(i, 1, line){
		for (auto col : v[i]){
			while (!s.empty()) s.pop();
			for (auto edge : e[col]){
				int x = edge.x, y = edge.y;
				int fx = gf(x), fy = gf(y);
				if (fx != fy){
					s.push(MP(fx, father[fx]));
					father[fx] = fy;
				}
			}

			cnt = 0;
			for (auto u : mp[col]) c[++cnt] = u.fi;
			rep(j, 1, cnt - 1){
				rep(k, j + 1, cnt){
					int x = c[j], y = c[k];
					if (x > y) swap(x, y);
					if (gf(x) == gf(y)) ++ret[MP(x, y)];
				}
			}

			while (!s.empty()){
				father[s.top().fi] = s.top().se;
				s.pop();
			}

		}
	}

	rep(i, line + 1, n){
		for (auto col : v[i]){
			memset(father, 0, sizeof father);
			for (auto edge : e[col]){
				int x = edge.x, y = edge.y;
				int fx = getfather(x), fy = getfather(y);
				if (fx != fy) father[fx] = fy;
			}

			rep(j, 1, qu) if (getfather(q[j].x) == getfather(q[j].y)) ++q[j].ans;
		}
	}

	rep(i, 1, qu) printf("%d\n", q[i].ans + ret[MP(q[i].x, q[i].y)]);
	return 0;
}