PAT | T1016 Uniqueness of MST

测试点3答案错误，31分

#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_set>

using namespace std;

const int maxn = 510;
typedef long long ll;

struct Edge{
	int u,v;
	ll w; // 权值
};

vector<Edge> edges;
int pre[maxn];

bool cmp(Edge a,Edge b){
	return a.w < b.w;
}

int unionSearch(int root){ // 查找根节点并路径压缩
	int son,tmp;
	son = root;
	while(root != pre[root]) root = pre[root]; // 查找根节点
	while(son != root){ // 路径压缩
		tmp = pre[son];
		pre[son] = root;
		son = tmp;
	}
	return root;
}

void join(int root1,int root2,int &cnt){ // 合并
	int x,y;
	x = unionSearch(root1);
	y = unionSearch(root2);
	if(x != y) pre[x] = y,cnt--;
}

void canJoin(Edge a,Edge b,bool &unique,ll &res,int &cnt){ // 看这两条边加到最小生成树中是否效果一样，如果一样说明结果不唯一
	int x1 = unionSearch(a.u),x2 = unionSearch(b.u);
	int y1 = unionSearch(a.v),y2 = unionSearch(b.v);
	if(x1 == y1) return; // 这条边不加到最小生成树中
	if(((x1 == x2 && y1 == y2) || (x1 == y2 && x2 == y1)) && a.w == b.w) unique = false; // 这两条边效果相同
	join(a.u,a.v,cnt);
	res += a.w;
	return;
}

int main(){
	int n,m;
	cin>>n>>m;
	int cnt = n;
	for(int i = 1;i <= n;i++) pre[i] = i; // 初始化pre数组
	for(int i = 0;i < m;i++){
		Edge tmp;
		cin>>tmp.u>>tmp.v>>tmp.w;
		edges.push_back(tmp);
		join(tmp.u,tmp.v,cnt);
	}
	if(cnt > 1){
		printf("No MST\n%d",cnt);
		system("pause");
		return 0;
	}
	sort(edges.begin(),edges.end(),cmp);
	for(int i = 1;i <= n;i++) pre[i] = i; // 初始化pre数组
	int siz = edges.size();
	cnt = n;
	ll res = 0;
	bool unique = true;
	for(int i = 0;cnt > 1 && i < siz;i++){
		if(i != siz - 1) canJoin(edges[i],edges[i + 1],unique,res,cnt);
	}
	printf("%lld\n",res);
	if(unique) printf("Yes");
	else printf("No");
	system("pause");
	return 0;
}