Little W and Contest（并查集）

Little W and Contest（并查集）

题目传送门

Little W and Contest

题意

俱乐部有n个人，其中有标号为1和2的两种人， 一个队伍至少保包含两个标号为2的人
最开始大家相互不熟悉，每天会有人相互认识，也可以通过认识的人去认识其他人（并查集）
要求队伍中的人互相不认识，求每天可以组成的队伍

思路

第一天大家都不认识，我们可以选择2 2 1和2 2 2，求出此时的ans
后面的几天，可以将人分成三种，当天认识的u，v和其他的w，因为u和v肯定是不互相认识的，所以认识u的，认识v的，和其他的人w，很明显当有人认识后就会出现需要除去的方案，分别从u，v，w中选
也就是这四种情况

    ans-=p2[u]*p2[v]*(cnt2-p2[u]-p2[v]);    //除去的方案为2 2 2
    ans-=p2[u]*p2[v]*(cnt1-p1[u]-p1[v]);    //除去的方案为2 2 1
    ans-=p1[u]*p2[v]*(cnt2-p2[u]-p2[v]);    //除去的方案为1 2 2
    ans-=p2[u]*p1[v]*(cnt2-p2[u]-p2[v]);    //除去的方案为2 1 2

其中的p1和p2是以i为根节点的集合中的数量
操作完后记得合并u和v的集合，同时维护根节点中的数量

代码

#include<cstdio>
#include<iostream>
#include<cmath>
using namespace std;
const double pi = acos(-1.0);
#define INF 0x3f3f3f3f
// #define TDS_ACM_LOCAL
typedef long long ll;
const int mod=1e9 +7;
const int N=1e5 +9;
ll f[N], p1[N], p2[N], a;
ll n, cnt1, cnt2, u, v;
ll ans;

void work(){
    ans-=p2[u]*p2[v]*(cnt2-p2[u]-p2[v]);    //除去的方案为2 2 2
    ans-=p2[u]*p2[v]*(cnt1-p1[u]-p1[v]);    //除去的方案为2 2 1
    ans-=p1[u]*p2[v]*(cnt2-p2[u]-p2[v]);    //除去的方案为1 2 2
    ans-=p2[u]*p1[v]*(cnt2-p2[u]-p2[v]);    //除去的方案为2 1 2
    cout<<ans%mod<<endl;
}

void Union(){       
    f[u]=v;                         //将u和v所处的集合合并
    p1[v]+=p1[u], p2[v]+=p2[u];     //集合含有的人合并
}

int find(int x){
    if(x==f[x]) return x;
    return f[x]=find(f[x]);
}

void solve(){
    cin>>n;
    ans=cnt1=cnt2=0;
    for(int i=1; i<=n; i++){
        cin>>a;
        if(a==1) cnt1++, p1[i]=1, p2[i]=0;
        else    cnt2++, p2[i]=1, p1[i]=0;
        f[i]=i;
    }
    if(cnt2>=1) ans+=cnt2*(cnt2-1)*cnt1/2;       //2 2 1的情况
    if(cnt2>=2) ans+=cnt2*(cnt2-1)*(cnt2-2)/6;  //2 2 2的情况
    cout<<ans%mod<<endl;
    for(int i=1; i<n; i++){
        cin>>u>>v;
        u=find(u), v=find(v);
        work();
        Union();
    }
    return ;
}   

int main(){
	ios::sync_with_stdio(0);
	cin.tie(0), cout.tie(0);
#ifdef TDS_ACM_LOCAL
    freopen("D:\\VS code\\.vscode\\testall\\in.txt", "r", stdin);
    freopen("D:\\VS code\\.vscode\\testall\\out.txt", "w", stdout);
#endif
    int T;
    cin>>T;
    while(T--)  solve();
    return 0;
}