【PAT甲级】1034. Head of a Gang (30)

#include <stdio.h>
#include <iostream>
#include <map>
#include <string.h>
#include <string>
#include <deque>
#include <algorithm>
using namespace std;

typedef struct Gang {
int head_id;
int num;
public:
Gang(int a, int b) {
head_id = a;
num = b;
}
} Gang;

deque<Gang> gangs;
map<string, int> m;
map<int, string> r;
int s[2001];
int w[2001];
int w_total[2001];
int mem[2001];
int n, k;
int cnt;

bool com(Gang a, Gang b) {
if (r.find(a.head_id)->second < r.find(b.head_id)->second)
return true;
return false;
}

int main(int argc, char *argv[]) {
scanf("%d %d", &n, &k);
int i, j;
for (i = 0; i < 2001; i++) {
w[i] = 0;
s[i] = -1;
w_total[i] = 0;
mem[i] = 0;
}
cnt = 0;
for (i = 0; i < n; i++) {
string a, b;
int ai, bi;
int t;
cin >> a >> b >> t;
if (m.find(a) != m.end()) {
ai = m.find(a)->second;
w[ai] += t;
} else {
ai = cnt++;
m.insert(make_pair(a, ai));
r.insert(make_pair(ai, a));
w[ai] = t;
}
if (m.find(b) != m.end()) {
bi = m.find(b)->second;
w[bi] += t;
} else {
bi = cnt++;
m.insert(make_pair(b, bi));
r.insert(make_pair(bi, b));
w[bi] = t;
}
int pa = ai;
int pb = bi;
while (s[pa] >= 0) {
pa = s[pa];
}
while (s[pb] >= 0) {
pb = s[pb];
}
if (pa != pb) {//两者不在一个集合，则加入同一集合中；否则什么都不做
s[pa] += s[pb];
s[pb] = pa;
}
}

for (i = 0; i < cnt; i++) {
if (s[i] >= 0 && s[s[i]] >= 0) {
int p = i;
while (s[p] >= 0) {
p = s[p];
}
s[i] = p;
}
}

for (i = 0; i < cnt; i++) {
if (s[i] >= 0) {//////
w_total[s[i]] += w[i];
mem[s[i]]++;
}
else {
w_total[i] += w[i];
mem[i]++;
}
}
for (i = 0; i < cnt; i++) {
if (s[i] < 0 && w_total[i] > 2 * k) {
int head_id = i;
int maxw = 0;
for (j = 0; j < cnt; j++) {
if ((s[j] == i || j == i) && w[j] > maxw) {
head_id = j;
maxw = w[j];
}
}
if (mem[i] > 2) {//gang - more than two persons
Gang g(head_id, mem[i]);
gangs.push_back(g);
}
}
}
sort(gangs.begin(), gangs.end(), com);
printf("%d\n", int(gangs.size()));
for (i = 0; i < int(gangs.size()); i++) {
cout << r.find(gangs[i].head_id)->second << " " << gangs[i].num << endl;
}
return 0;
}