题解 | #Head of a Gang#

并查集思想

本题也可以使用并查集解决。在使用并查集时，只要注意合并函数中需要总是保持点权更大的结点为集合的根结点(原先的合并函数是随意指定其中一个根结点为合并后集合的根结点)，就能符合题目的要求。

而为了达到题目对总边权与成员人数的要求，需要定义两个数组:

一个数组用来存放以当前结点为根结点的集合的总边权（totalWeight）。

另一个数组用来存放以当前结点为根结点的集合中的成员人数（gangs）。

这样当所有通话记录合并处理完毕后，这两个数组就自动存放了每个集合的总边权和成员人数，再根据题意进行筛选即可。 （摘自算法笔记）

#include <iostream>
#include <string>
#include <vector>
#include <map>
using namespace std;

const int MAXN = 2000;

map<string, int> mp1;    // name ---> index
map<int, string> mp2;    //index ----> name

vector<int> s(MAXN, -1);     //并查集数组
vector<int> weight(MAXN, 0);    //存放每个顶点的点权值
vector<int> totalWeight(MAXN, 0);  //顶点为根节点的集合的总边权值

vector<int> gangs[MAXN];      //gangs[root] 存放以root为根的该集合的所有元素
map<string, int> result;      //head ---> 人数

int idx = 0;

int stringToInt(string s) {
    if (mp1.find(s) != mp1.end()) {
        return mp1[s];
    }
    else {
        mp1[s] = idx;
        mp2[idx] = s;
        return idx++;
    }
}

int Find(vector<int>& s, int x) {
    int root = x;
    while (s[root] >= 0) {
        root = s[root];
    }
    while (x != root) {
        int t = s[x];
        s[x] = root;
        x = t;
    }
    return root;
}

void Union(vector<int>& s, int a, int b, int w) {
    int roota = Find(s, a);
    int rootb = Find(s, b);
    if (roota == rootb) {
        totalWeight[roota] += w;   //更新顶点为根节点的集合的总边权值
        return;
    }
    if (s[roota] <= s[rootb]) {
        s[roota] += s[rootb];
        s[rootb] = roota;
        totalWeight[roota] += totalWeight[rootb] + w;
    }
    else {
        s[rootb] += s[roota];
        s[roota] = rootb;
        totalWeight[rootb] += totalWeight[roota] + w;
    }    
}

void Initial() {        //输入第二组测试数据时，恢复全局变量状态
    for (int i = 0; i < MAXN; i++) {
        s[i] = -1;
        weight[i] = 0;
        totalWeight[i] = 0;
        gangs[i].clear();
        result.clear();
        mp1.clear();
        mp2.clear();
    }
}

int main()
{
    int N, K;
    while (cin >> N >> K) {
        Initial();
        string s1, s2;
        int w;
        while (N--) {
            cin >> s1 >> s2 >> w;
            int u = stringToInt(s1);
            int v = stringToInt(s2);
            Union(s, u, v, w);
            weight[u] += w;           //更新顶点的点权值
            weight[v] += w;
        }
        for (int i = 0; i < idx; i++) {     //将每个人加入自己所在的集合，gangs[root]数组
            int id = Find(s, i);
            gangs[id].push_back(i);
        }
        int cnt = 0;
        for (int i = 0; i < idx; i++) {
            if (totalWeight[i] <= K || gangs[i].size() <= 2) continue;
            int h = i;
            for (auto p : gangs[i]) {       //找到该集合中点权值最大的人的index
                if (weight[p] > weight[h]) {
                    h = p;
                }
            }
            result[mp2[h]] = gangs[Find(s, h)].size();   //gangs[root]中存着h所在集合的人的总数
            cnt++;
        }
        cout << cnt << endl;
        for (auto it = result.begin(); it != result.end(); it++) {
            cout << it->first << " " << it->second << endl;
        }
    }
    return 0;
}

Union函数修改

void Union(vector<int>& s, int a, int b, int w) {
    int roota = Find(s, a);
    int rootb = Find(s, b);
    if (roota == rootb) {
        totalWeight[roota] += w;   //更新顶点为根节点的集合的总边权值
        return;
    }
    if (weight[roota] >= weight[rootb]) {
        s[rootb] = roota;
        totalWeight[roota] += totalWeight[rootb] + w;
    }
    else {
        s[roota] = rootb;
        totalWeight[rootb] += totalWeight[roota] + w;
    }    
}