Little W and Contest HDU

Parctice link: https://vjudge.net/problem/HDU-6795

Problem Description

There are n members in our ACM club. Little W wants to select three persons from our club to form a new team taking part in provincial ACM contests, as it is known by all of us that any ACM contest requires a normal team to have three members.
Little W has divided our club members into two role groups. The first group contains only readers who dedicate themselves to reading problems during contests, though sometimes they may also prepare drinking and food for the team. For the sake of measurement, we define the power of a reader as 1. The second part contains only coders who code and test programs all the time, and similarly, we define the power of a coder as 2.

Little W thinks it will be a tremendous disaster when a team has two readers because in that case, the total power of this team is less than 5 and thus it has a high risk to fail the contest. To avoid that, Little W thinks a new team must have at least two coders.

Additionally, Little W defines the relationship between club members with transitivity. That is, for every three members A, B, and C, if A is familiar with B, and B is familiar with C, then A will be familiar with C through B instantly. Based on the definition, it is forbidden for the team to have any two members familiar with each other.

At first, no member of our club is familiar with any other, and then Little W will repeatedly make an introduction between two members who are currently strangers to each other until each member is familiar with all the others. During this process, there will be exactly (n−1) introductions.

Now, for i=1,2,…,n, Little W wants you to count the combinations of three club members that can form a new team after the first (i−1) introductions have been made. However, the numbers of combinations may be quite gigantic, so you just need to report each number in modulo (109+7).

 Input

There are several test cases.

The first line contains an integer T (1≤T≤10), denoting the number of test cases. Then follow all the test cases.

For each test case, the first line contains an integer n (1≤n≤105), denoting the number of members in this club.

The second line contains n integers consisting of only 1 and 2, where the i-th integer represents the power of the i-th member.

The next (n−1) lines describe all introductions in chronological order of occurrence, where each line contains two integers u and v (1≤u,v≤n,u≠v), representing an introduction between the u-th member and the v-th member, who are currently strangers to each other.

It is guaranteed that the sum of n is no larger than 106.

 Output

For each test case, output n lines, where the i-th line contains an integer, denoting the number of combinations of three club members, in modulo (109+7), that can form a new team after the first (i−1) introductions have been made.

 

题意：给你 n 个人，这 n 个人有 1 或者 2 的权值，再给你分别给你 n - 1条边，你需要在这 n 个人中挑选 3 个人，且这三个人的权值总和大于等于5，并且这三个人在图上是不连通的，问你每添加一条边，可以挑选的总方案数分别是多少。

思路：首先，考虑没有边时，我们假设 cnt1 是拥有权值为 1 的人的数量，cnt2 是拥有权值 2 的人的数量，那么初始方案就是

$ans$

           接下来就要考虑有边的情况，首先 设 p1[ i ] 是以 i 为根节点的连通块的 1 的权值的人的数量， p2[ i ]就是以 i 为根节点的连通块的 2 的权值的人的数量。然后引入一条边（u，v）时，设三个集合 G( u )即以u为根节点的连通块，G( v )时以 v 为根节点的连通块，G( r )是其他连通块，那么引入这条边后，就有一些方案不能成立，就是三个人分别属于G( u)、G( v )和G( r )的情况。也就是（2、2、1）、（2、2、2）、（1、2、2）和（2、1、2）四种情况，依次减去这四种情况，就是在添加这条边后的方案数，最后用并查集合并两个点所在的连通块。

代码:

#include<bits/stdc++.h>
#define ll long long
#define INF 0x3f3f3f3f
#define mem(a,x) memset(a,x,sizeof(a))  
#define _for(i,a,b) for(int i=a; i< b; i++)
#define _rep(i,a,b) for(int i=a; i<=b; i++)
#define ios ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
using namespace std;
const int MOD = 1e9+7 ;
const int maxn = 100005 ;

int fa[maxn];
int p1[maxn],p2[maxn];   //以i为根的连通块含1 和含2 
int a[maxn];
int n;
ll ans=0;
int find(int x)
{
  return (fa[x]==x)?x:(fa[x]=find(fa[x]));
}

int main()
{
    int T;
    scanf("%d",&T);
    while(T--){
      ll x=0,y=0;
      ans=0;
      scanf("%d",&n);
      for(int i=1;i<=n;i++){
        scanf("%d",&a[i]);
        if(a[i]==1){
          x++;
          p1[i]=1;
          p2[i]=0;
        }
        if(a[i]==2){
          y++;
          p1[i]=0;
          p2[i]=1;
        }
        fa[i]=i;
      }
      ans=(ll)(y*(y-1)/2)*x+(ll)(y*(y-1)*(y-2))/6;
      ans=ans%MOD;
      printf("%lld
",ans);
      for(int i=1;i<n;i++){
          int u,v;
          ll k=0;
          scanf("%d %d",&u,&v);
          int pu=find(u),pv=find(v);
          k=(k+(ll)(p2[pu]*p2[pv])*(y-p2[pu]-p2[pv]))%MOD;
          k=(k+(ll)(p2[pu]*p1[pv])*(y-p2[pu]-p2[pv]))%MOD;
          k=(k+(ll)(p1[pu]*p2[pv])*(y-p2[pu]-p2[pv]))%MOD;
          k=(k+(ll)(p2[pu]*p2[pv])*(x-p1[pu]-p1[pv]))%MOD;
          ans=(ans-k+MOD)%MOD;
          printf("%lld
",ans);
          fa[pv]=pu;
          p1[pu]+=p1[pv];
          p2[pu]+=p2[pv];
      }

    }
    return 0;
}