Little W and Contest(并查集)
题意: acm组队,有n名队员,每名成员担当读题或者编码的任务,读题能力为1,编程能力为2,每个队伍三名成员,成员能力值至少为5,刚开始队员之间相互不熟悉,每过一天就会有人相互熟悉,并且也会熟悉双方的朋友,问每天最多可以组多少队。
hint:用并查集。记录并查集中能力为1的人数,能力为2的人数,当两个并查集合并时要去掉两个集合之前能配对的组合数,并且合并两个并查集的能力值个数。
#include <bits/stdc++.h>
using namespace std;
#define LL long long
#define sc(a) scanf("%d", &a)
#define s2(a, b) scanf("%d%d", &a, &b)
#define ss(a) scanf("%s", a)
#define mem(a, b) memset(a, b, sizeof(a))
const int N = 1e5 + 5;
const int mod = 1e9 + 7;
int pre[N], cnt1[N], cnt2[N], a[N];
int find(int a)
{
if (a == pre[a])
return a;
return pre[a] = find(pre[a]);
}
LL a1, a2;
//a1表示能力值为1的总人数
//a2表示能力值为2的总人数
int main()
{
int t;
sc(t);
while (t--)
{
a1 = a2 = 0;
int n;
sc(n);
for (int i = 1; i <= n; i++)
{
sc(a[i]);
pre[i] = i;
if (a[i] == 1)
{
cnt1[i] = 1;
cnt2[i] = 0;
a1++;
}
else
{
cnt1[i] = 0;
cnt2[i] = 1;
a2++;
}
}
//刚开始可以组队的人数
LL ans = a2 * (a2 - 1) * (a2 - 2) / 6 + a2 * (a2 - 1) * a1 / 2;
cout << ans % mod << endl;
for (int i = 1; i < n; i++)
{
int x, y;
s2(x, y);
int u, v;
u = find(x), v = find(y);
//2 2 2 的组合 两个要合并的集合+其他集合
ans -= cnt2[u] * cnt2[v] * (a2 - cnt2[u] - cnt2[v]);
//2 2 1 的组合
ans -= cnt2[u] * cnt2[v] * (a1 - cnt1[u] - cnt1[v]);
//2 1 2 的组合
ans -= cnt2[u] * cnt1[v] * (a2 - cnt2[u] - cnt2[v]);
//1 2 2 的组合
ans -= cnt1[u] * cnt2[v] * (a2 - cnt2[u] - cnt2[v]);
//合并,人数也要合并
pre[u] = v;
cnt1[v] += cnt1[u];
cnt2[v] += cnt2[u];
cout << ans % mod << endl;
}
}
system("pause");
return 0;
}