2020多校第三场Little W and Contest

思路

用并查集，维护三种连通快中。先求出第一天的ans，每次从不同块中去取不满足的情况，存在四种可能，每次用前一次的减去不满足的情况个数，得到最终ans。

代码

#include 
#define hi cout<<"hi\n";
#define inf 0x3f3f3f3f
#define ll long long
#define ull unsigned long long
using namespace std;
const ll mod = 1e9+7;
const int N = 100005;
ll power(long long a, long long b, long long mode)
{
    long long sum = 1;
    while (b)
    {
        if (b & 1)
        {
            sum = (sum * a) % mode;
            b--;
        }
        b /= 2;
        a = a * a % mode;
    }
    return sum;
}
ll gcd(int a,int b)
{
    if(b==0)
        return a;
    else
        return gcd(b,a%b);
}
ll lcm(int a,int b)
{
    return a/gcd(a,b)*b;
}
int tcase,n,i,n1,n2,k,ans,j=1,u,v,pu,pv,x;
int p1[N],p2[N],p[N];
int Find(int xx)//并查集find函数
{
    if(p[xx]!=xx)
	p[xx]=Find(p[xx]);
    return p[xx];
}
int main(){
	scanf("%d",&tcase);
	while(tcase--){
		scanf("%d",&n);
		n1 = n2 = 0;
		for(i=1;i<=n;i++){
			scanf("%d",&x); p[i]=i;
			if(x==1){
				n1++;p1[i]=1;p2[i]=0;
			}
			else{
				n2++;p2[i]=1;p1[i]=0;
			}
		}

		ans = ((ll)n2*(n2-1)*(n2-2)/6%mod + (ll)n1*n2%mod*(n2-1)/2%mod)%mod;
		cout<<ans<<endl;
		for(i=1;i<n;i++){
		        k=0;
			    scanf("%d%d",&u,&v);
		        int pu=Find(u), pv=Find(v);

				k=(k+(ll)p1[pu]*p2[pv]*(n2-p2[pu]-p2[pv]))%mod;

                k=(k+(ll)p2[pu]*p1[pv]*(n2-p2[pu]-p2[pv]))%mod;

                k=(k+(ll)p2[pu]*p2[pv]*(n2-p2[pu]-p2[pv]))%mod;

                k=(k+(ll)p2[pu]*p2[pv]*(n1-p1[pu]-p1[pv]))%mod;

                ans = (ans - k + mod)%mod;
				p[pv]=pu;
                p1[pu]+=p1[pv]; p2[pu]+=p2[pv];
                p1[pv]=0; p2[pv]=0;
                cout<<ans<<endl;
	      }

  }
	return 0;
}