2020年HDU多校第三场 1005 Little W and Contest（并查集与数学）

2020年HDU多校第三场 1005 Little W and Contest（并查集与数学）

题意：给n个人与每个的能力值（2或1），组一个3人的队能力和至少5以上，同在一个集合的人不能组队，最开始每个人都在自己的集合，询问n-1次，每次将任意两人所在的集合合并之后输出多少种组队方案。

题解：先撇开5以上不说，解决3人组队的问题，最开始就一个comb(n,3)，主要是合并以后怎么求，我们可以考虑合并以后对原结果的影响，设x集合的人，与y集合的人合并，那么原本x的人，与y的人的组合方案我就得去掉，去掉的值为x * y * (n-x-y)，即x中取一人，y中取一人，x，y之外取一个人的方案数，那么为什么不用去除x取一人，y中取二人的方案呢，因为在y之前合并的时候已经去除的此方案数，现在引入题目为5以上的要求，枚举一下就好了，221，222，122，212;

#include<iostream>
using namespace std;
#define ll long long
ll t,lin,n,fa[100007],sum2,sum1,ans,u,v,res;
struct madoka{
    ll one;
    ll two;
    ll siz;
}ma[100007];
const long long mod = 1e9+7;
long long fac[2000006];
long long qpow(long long x, long long n) {
    long long res = 1;
    for (; n; n >>= 1, x = x * x % mod)
        if (n & 1) res = res * x % mod;
    return res;
}
long long inv(long long a) {
    return qpow(a, mod-2)%mod;
}
void solve() {
    fac[0] = 1;
    for(int i = 1;i <= 2000006; i++) {
        fac[i] = (fac[i-1]*i)%mod;
    }
}
long long comb(long long n, long long k) {
    if(k > n) return 0;
    return (fac[n]*inv(fac[k])%mod * inv(fac[n-k])%mod);
}

ll fin(int p){
    if(p==fa[p])return p;
    else{
        return fa[p]=fin(fa[p]);
    }
}
void go(int f1,int f2){
    ll o1=ma[f1].one;
    ll o2=ma[f2].one;
    ll t1=ma[f1].two;
    ll t2=ma[f2].two;
    res=(res+comb(t1,1)*comb(t2,1)%mod*comb(sum2-t1-t2,1))%mod;
    res=(res+comb(t1,1)*comb(t2,1)%mod*comb(sum1-o1-o2,1))%mod;
    res=(res+comb(o1,1)*comb(t2,1)%mod*comb(sum2-t1-t2,1))%mod;
    res=(res+comb(t1,1)*comb(o2,1)%mod*comb(sum2-t1-t2,1))%mod;
    ma[f1].one+=ma[f2].one;
    ma[f1].two+=ma[f2].two;
    ma[f1].siz+=ma[f2].siz;
    fa[f2]=f1;
}
void init(){
    for(int i=1;i<=n;i++){
        fa[i]=i;
        ma[i].one=0;
        ma[i].two=0;
        ma[i].siz=0;
    }
    sum1=0;
    sum2=0;
    ans=0;
    res=0;
}
int main(){
    solve();
    scanf("%lld",&t);
    while(t--){
        scanf("%lld",&n);
        init();
        for(int i=1;i<=n;i++){
            scanf("%lld",&lin);
            if(lin==1){
                ma[i].one++;
                sum1++;
            }
            else {
                ma[i].two++;
                sum2++;
            }
            ma[i].siz=1;
        }
        ans=(comb(sum1,1)*comb(sum2,2)%mod+comb(sum2,3))%mod;
        printf("%lld\n",ans);
        for(int i=1;i<n;i++){
            scanf("%lld%lld",&u,&v);
            int f1=fin(u),f2=fin(v);
            go(f1,f2);
            printf("%lld\n",(ans-res+mod)%mod);
        }
    }
}