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