PTA 1034. Head of a Gang (30)

使用并查集

#include<stdio.h>
#include<string.h>
#include<string>
#include<iostream>
#include<queue>
#include<stack>
#include<math.h>
#include<string>
#include<algorithm>
#include<map>
using namespace std;
vector<int>v[2];	//保存输入的对话 
int a[260000];	//并查集 
int value[260000];	//每个人的weight 
int all[260000];	//每个连通块的总weight 使用时除2 
 
int find(int x){
	if(a[x] < 0)	return x;
	return a[x] = find(a[x]);
}

void insert(int x,int y){
	int p = find(x);
	int q = find(y);
	if(p == q)	return ;
	if(value[p] > value[q])	a[q] = p;	//保证连通块头点为最大weight 
	else a[p] = q;
}

int main(){
	int n,m;
	scanf("%d%d",&n,&m);
	memset(value,0,sizeof(value));
	memset(all,0,sizeof(all));
	memset(a,-1,sizeof(a));
	for(int i=0;i<n;i++){	//数字化，每一个字母对应两位数，直接用数组 
		char x[4],y[4];
		int p,q,t;
		scanf("%s",x);
		getchar();
		scanf("%s",y);
		scanf("%d",&t);
		p = (x[0]-'A')*10000+(x[1]-'A')*100+(x[2]-'A');
		q = (y[0]-'A')*10000+(y[1]-'A')*100+(y[2]-'A');
		value[p] += t;
		value[q] += t;
		v[0].push_back(p);
		v[1].push_back(q);
	}
	for(int i=0;i<n;i++){	//并查集 
		insert(v[0][i],v[1][i]);
	}
	bool book[260000];//用来标记该点是否已经计算过 
	memset(book,false,sizeof(book));
	vector<int>r;//保存所有连通块的头结点 
	for(int i=0;i<n;i++){
		for(int j=0;j<2;j++){
			if(!book[v[j][i]]){
				book[v[j][i]] = 1;
				if(a[v[j][i]] >= 0){//不为头结点 
					int point = find(v[j][i]);//找到头节点 
					a[point]--;//用负值来保存该连通块共有几个点，初始值为-1 
					all[point] += value[v[j][i]];//合计连通块weight 
				}
				else{//头结点 
					all[v[j][i]] += value[v[j][i]];
					r.push_back(v[j][i]);
				}
			}
		}
	}
	vector<int >real;//保存结果 
	for(int i=0;i<r.size();i++){
		if(a[r[i]] < -2 && all[r[i]]/2 > m)	real.push_back(r[i]);	//判断是否符合条件 
	}
	printf("%d\n",real.size());
	sort(real.begin(),real.end());
	for(int i=0;i<real.size();i++){//复原 
		int rea = real[i];
		char temp[4];
		temp[0] = rea/10000+'A';
		rea%=10000;
		temp[1] = rea/100+'A';
		rea%=100;
		temp[2] = rea+'A';
		temp[3] = 0;
		printf("%s %d\n",temp,-a[real[i]]);
	}
	return 0;
}