PAT A1034 Head Of Gang

用并查集分割团伙，判断输出~

#include<bits/stdc++.h>
using namespace std;
const int maxn=10010;
int father[maxn],isRoot[maxn]={0},weight[maxn];
unordered_map<string,int> pos1;
unordered_map<int,string> pos2;
unordered_map<string,int> ans;
struct node {
    string id;
    int total=0;
    int num=0;
}Node[maxn];
struct gang {
    string id;
    int num;
};
void init () {
    for (int i=0;i<maxn;i++) 
    father[i]=i;
}
int findfather (int x) {
    int a=x;
    while (x!=father[x]) 
    x=father[x];
    while (a!=father[a]) {
        int z=a;
        a=father[a];
        father[z]=x;
    }
    return x;
}
void Union (int a,int b) {
    int faA=findfather(a);
    int faB=findfather(b);
    if (faA!=faB) {
        if (weight[faA]>weight[faB]) father[faB]=faA;
        else father[faA]=faB;
    }
}
bool cmp (gang a,gang b) {
    return a.id<b.id;
}
int main () {
    int N,K;
    scanf ("%d %d",&N,&K);
    init ();
    string s1,s2;
    int cnt=1,x;
    vector<pair<string,string>> v1;
    for (int i=0;i<N;i++) {
        cin>>s1>>s2>>x;
        if (pos1[s1]==0) {
            pos1[s1]=cnt;
            pos2[cnt++]=s1;
        }
        if (pos1[s2]==0) {
            pos1[s2]=cnt;
            pos2[cnt++]=s2;
        }
        weight[pos1[s1]]+=x;
        weight[pos1[s2]]+=x;
        v1.push_back({s1,s2});
    }
    for (int i=0;i<v1.size();i++)
    Union (pos1[v1[i].first],pos1[v1[i].second]);
    for (int i=1;i<cnt;i++) {
        Node[findfather(i)].total+=weight[i];
        Node[findfather(i)].num++;
    }
    for (int i=1;i<cnt;i++) {
        if (Node[i].total>K*2&&Node[i].num>2) 
        ans[pos2[i]]=Node[i].num;
    }
    printf ("%d\n",ans.size());
    vector<gang> vi;
    for (auto it=ans.begin();it!=ans.end();it++)
    vi.push_back({it->first,it->second});
    sort (vi.begin(),vi.end(),cmp);
    for (int i=0;i<vi.size();i++)
    cout<<vi[i].id<<" "<<vi[i].num<<endl;
    return 0;
}