HDU-6795 Little W and Contest 【2020 Multi-University Training Contest 3】【并查集】

题目

HDU-6795

题意

n个人，每个人的能力值为1或者2，现在进行组队，组队规则如下：
1、三个能力值为2可以组一队；
2、两个能力值为2和一个能力值为1的可以组一队；
3、三人必须相互不认识才能组队；
起初所有人都不认识对方，接下来小W给其中两个人介绍，使他们相互认识（a认识b，b认识c，c也就认识a），小W会进行n-1次介绍；
问一开始这n个人可以组多少不同的队伍，以及小W每次介绍后，这n个人又可以组多少不同的队伍；

题解

首先对于这种交友问题明显是要使用并查集，大致思路就是我们先计算出一开始的组队数，当加入一对新的朋友时查看两个集合中1和2的数量，再减去两人交友后影响的组队数；

用num[1] 和num[2] 分别记录全部的1和2的数量；
用cnt[i][1] 和cnt[i][2] 分别记录每个集合中1和2的数量；
一开始的队伍数：

ans = num[2] * (num[2] - 1) / 2 * num[1] + num[2] * (num[2] - 1) * (num[2] - 2) / 6;

两个人交友后，需要将两个集合中cnt的值合并，然后减去原先两个集合中的人和其他集合的人组队的数量（即原先两个集合的人可以各取一人再加上其他集合的一人组为一队，现在需要去除，因为这两个集合已经合并成一个集合，不能在一个队伍）；计算公式如下：

u1为集合u能力值为1的人数，u2为集合u能力值为2的人数
v1为集合v能力值为1的人数，v2为集合v能力值为2的人数
uv1为集合u和v能力值为1的人数之和，uv2为集合u和v能力值为2的人数之和

ans -= u1 * v2 * (num[2] - uv2);
ans -= u2 * v1 * (num[2] - uv2);
ans -= u2 * v2 * (n - uv2 - uv1);

ps:注意取模

代码

/*
 * @author: arc
 * @date: 2020-08-02 15:57:58
 */
int r[maxn];
int cnt[maxn][3];
int num[3];

void init(int n){
    for(int i = 1; i <= n; i++){
        r[i] = i;
    }
}

int __find(int k){
    if(r[k] == k) return k;
    return r[k] = __find(r[k]);
}

int __merge(int a, int b){
    int fa = __find(a);
    int fb = __find(b);
    if(fa != fb) {
        r[fb] = fa;
        cnt[fa][1] += cnt[fb][1];
        cnt[fa][2] += cnt[fb][2];
    }
    return fa;
}

int main(){
    int cas;
    scanf("%d", &cas);
    while(cas--){
        int n, a;
        scanf("%d", &n);
        init(n);
        num[1] = 0;
        num[2] = 0;
        memset(cnt, 0, sizeof(cnt));
        for (int i = 1; i <= n; i++){
            scanf("%d", &a);
            cnt[i][a]++;
            num[a]++;
        }
        ll ans = 1LL * num[2] * (num[2] - 1) / 2LL % mod * num[1] % mod + 1LL * num[2] * (num[2] - 1) * (num[2] - 2) / 6LL % mod;
        printf("%lld\n", ans % mod);
        int u, v;
        for (int i = 1; i < n; i++){
            scanf("%d%d", &u, &v);
            if(i >= n-2){
                printf("0\n");
                continue;
            }
            int fu = __find(u);
            int fv = __find(v);
            int u1 = cnt[fu][1];
            int u2 = cnt[fu][2];
            int v1 = cnt[fv][1];
            int v2 = cnt[fv][2];
            int fuv = __merge(u, v);
            int uv1 = cnt[fuv][1];
            int uv2 = cnt[fuv][2];
            ans -= 1LL * u1 * v2 % mod * (num[2] - uv2) % mod;
            if(ans < 0)
                ans += mod;
            ans -= 1LL * u2 * v1 % mod * (num[2] - uv2) % mod;
            if(ans < 0)
                ans += mod;
            ans -= 1LL * u2 * v2 % mod * (n - uv2 - uv1) % mod;
            if(ans < 0)
                ans += mod;
            printf("%lld\n", ans % mod);
        }
    }
}