2020杭电多校第三场1005（并查集）

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;
}