HDU 6795【2020杭电多校赛第三场1005】【并查集】【组合数学】

Little W and Contest

Sourse

HDU 6795 Little W and Contest

Solution

有两种人，读题的和写代码的，读题的人能力值是1，写代码的人能力值是2，组成一个团队最少需要能力5，所以至少要有2个写代码的人才能组成一个团队。
团队里的人，任意两个人不能互相认识。每个回合教练都会介绍任意两个人认识（因此这两人不能在同一个团队里），并且认识的人之间具有传递性，即题目说明了a认识b，说明了b认识c，a就会视作认识c。问每个回合过后有多少种组成团队的方式。
这种点相连的问题首先想到的是并查集，然后问组合方式必然是组合数学，但是比赛的时候只能想到知识点，没有太多思路，赛后看别人的题解才恍然大悟。
设能力2的人为ability[2]，能力1的人为ability[1]，朋友圈根节点为rt能力为x的人为friends[rt][x]。
则最开始组成团队的方式为：
ans = $C_{ability[2]}^3$ + $C_{ability[2]}^2$ × $C_{ability[1]}^1$
而u和v认识之后，有以下4种情况需要从现有ans中剔除。
（1）从u的圈中选出一个能力2的人，从v的圈中选出一个能力2的人，从剩下的人中选出一个能力为1的人。此情况下的组合数为：

friends[rtu][2] * friends[rtv][2] * (ability[1] - friends[rtu][1] - friends[rtv][1]);

（2）从u的圈中选出一个能力2的人，从v的圈中选出一个能力2的人，从剩下的人中选出一个能力为2的人。此情况下的组合数为：

friends[rtu][2] * friends[rtv][2] * (ability[2] - friends[rtu][2] - friends[rtv][2]);

（3）从u的圈中选出一个能力2的人，从v的圈中选出一个能力1的人，从剩下的人中选出一个能力为2的人。此情况下的组合数为：

friends[rtu][2] * friends[rtv][1] * (ability[2] - friends[rtu][2] - friends[rtv][2]);

（4）从u的圈中选出一个能力1的人，从v的圈中选出一个能力2的人，从剩下的人中选出一个能力为2的人。此情况下的组合数为：

friends[rtu][1] * friends[rtv][2] * (ability[2] - friends[rtu][2] - friends[rtv][2]);

用cin必须要关同步，不然会T，res和ans要开long long，friends[rt][x]相乘的时候需要先转换成long long，干脆全部开long long了。记得取模。我补题的的时候就是res最后忘了取模，wa了好多发都没发现，感谢队友帮我debug.

Code

/*
 * @Author: SolitaryOrz 
 * @Date: 2020-07-29 11:48:50 
 * @Last Modified by: SolitaryOrz
 * @Last Modified time: 2020-07-29 13:07:12
 */
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn = 1e5+10;
const int mod = 1e9+7;
#define fast ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);

ll arr[maxn],par[maxn],friends[maxn][5],ability[5],ans;

void init(int n) {
    for(int i=0;i<=n;i++) {
        par[i] = i;
        friends[i][1] = 0;
        friends[i][2] = 0;
    }
    ability[1] = 0;
    ability[2] = 0;
}

int findRoot(int x) {
    return par[x] == x ? x : par[x] = findRoot(par[x]);
}

void unionVertices(int x, int y) {
    x = findRoot(x);
	y = findRoot(y);
	if(x == y)
        return;
	else {
		par[x] = y;
		friends[y][1] += friends[x][1];
		friends[x][1] = 0;
		friends[y][2] += friends[x][2];
		friends[x][2] = 0;
	}
}

ll C(int n, int m) {
    if(n < m) return 0;
    ll res = 1;
    for(int i=n;i>n-m;i--) {
        res *= i;
    }
    for(int i=1;i<=m;i++) {
        res /= i;
    }
    return res%mod;
}

void solve(int u, int v) {
    ll res = 0;
    int rtu = findRoot(u), rtv = findRoot(v);
    res += (friends[rtu][2] * friends[rtv][2] % mod) * (ability[1] - friends[rtu][1] - friends[rtv][1]) % mod;
    res += (friends[rtu][2] * friends[rtv][2] % mod) * (ability[2] - friends[rtu][2] - friends[rtv][2]) % mod;
    res += (friends[rtu][2] * friends[rtv][1] % mod) * (ability[2] - friends[rtu][2] - friends[rtv][2]) % mod;
    res += (friends[rtu][1] * friends[rtv][2] % mod) * (ability[2] - friends[rtu][2] - friends[rtv][2]) % mod;
    ans = (ans - res % mod + mod) % mod;
    unionVertices(u,v);
}

int main() {
    fast;
    int t;
    cin >> t;
    while(t--) {
        int n;
        cin >> n;
        init(n);
        for(int i=1;i<=n;i++) {
            cin >> arr[i];
            ability[arr[i]]++;
            friends[i][arr[i]]++;
        }
        ans = ((C(ability[2],3) % mod) + (C(ability[2],2)  * C(ability[1],1) % mod)) % mod;
        cout << ans << endl;
        for(int i=1;i<n;i++) {
            int u,v;
            cin >> u >> v;
            solve(u,v);
            cout << ans << endl;
        }
    }
}