題目鏈接 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;
}

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