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 ((10^9)+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≤(10^5)), 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 (10^6)
.
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 (10^9+7)
, that can form a new team after the first (i−1) introductions have been made.
Sample Input
1
5
2 2 2 1 1
4 5
1 4
2 1
3 2
Sample Output
7
7
3
0
0
大意是所有人分为两个组 然后给了n-1条介绍信息 问每给一条信息（介绍两个人互相认识，注意认识关系具有传递性）后有多少种凑出来一个队的组合。三个人能成队当且仅当至少有两个人属于coder组且三个人互不认识。
首先看到这个题就能知道是并查集。正难则反，直接计算答案不太好搞，不如求出总的方案数，然后减去每次合并的贡献（负贡献）。一开始没想到这是带权并查集，后来看到别人的博客豁然开朗。额外开辟两个数组size1和size2，其中size1[i]表示并查集中以i为代表元的集合中的reader的个数，size2[i]表示并查集中以i为代表元的集合中的coder的个数。每次merge两个集合时，首先计算贡献。此时将所有点分成三部分：x所在的集合(S_x)，y所在的集合(S_y)以及剩下的点(S_{res})。受合并集合影响的总共有四种情况（只能选择两个二号点一个一号点或者三个二号点）：
1）从(S_x)拿一个二号点，从(S_y)拿一个二号点，从(S_{res})拿一个一号点。
2）从(S_x)拿一个二号点，从(S_y)拿一个二号点，从(S_{res})拿一个二号点。
3）从(S_x)拿一个一号点，从(S_y)拿一个二号点，从(S_{res})拿一个二号点。
4）从(S_x)拿一个二号点，从(S_y)拿一个一号点，从(S_{res})拿一个二号点。
对应代码就是：
tmp = (tmp + size2[x] * size2[y] * (reader - size1[x] - size1[y])) % p;
tmp = (tmp + size2[x] * size2[y] * (coder - size2[x] - size2[y])) % p;
tmp = (tmp + size1[x] * size2[y] * (coder - size2[x] - size2[y])) % p;
tmp = (tmp + size2[x] * size1[y] * (coder - size2[x] - size2[y])) % p;
其中tmp为总贡献，reader和coder为一号点和二号点的个数。注意这里的x和y是输入的两个点的所在集合的代表元！！！
此外很重要的是merge函数：
void merge(ll x, ll y)//并查集的merge操作 记得维护size数组
{
x = get(x), y = get(y);
size1[y] += size1[x];//先维护再并根
size2[y] += size2[x];
fa[x] = y;
}
以及组合数取模板子：
ll inv(ll a)
{
return a==1?1:(ll)(p-p/a)*inv(p%a)%p;
}
ll C(ll n,ll m)//组合数快速取模
{
if(m<0)return 0;
if(n<m)return 0;
if(m>n-m)m=n-m;
ll up=1,down=1;
ll i;
for(i=0;i<m;i++)
{
up=up*(n-i)%p;
down=down*(i+1)%p;
}
return up*inv(down)%p;
}
ll get(ll x)
{
if(x == fa[x]) return x;
return fa[x] = get(fa[x]);
}
以及别忘开long long！
代码如下：
#include <bits/stdc++.h>
#define ll long long
using namespace std;
const long long p = 1000000007;
ll n, fa[100005], size1[100005], size2[100005], reader, coder;
ll ans;
ll inv(ll a)
{
return a==1?1:(ll)(p-p/a)*inv(p%a)%p;
}
ll C(ll n,ll m)//组合数快速取模
{
if(m<0)return 0;
if(n<m)return 0;
if(m>n-m)m=n-m;
ll up=1,down=1;
ll i;
for(i=0;i<m;i++)
{
up=up*(n-i)%p;
down=down*(i+1)%p;
}
return up*inv(down)%p;
}
ll get(ll x)
{
if(x == fa[x]) return x;
return fa[x] = get(fa[x]);
}
void merge(ll x, ll y)//并查集的merge操作 记得维护size数组
{
x = get(x), y = get(y);
size1[y] += size1[x];//先维护再并根
size2[y] += size2[x];
fa[x] = y;
}
int main()
{
int t;
cin >> t;
while(t--)
{
cin >> n;
reader = coder = 0;
for(ll i = 1; i <= n; i++)
{
fa[i] = i;
ll temp;
scanf("%lld", &temp);
if(temp == 1) reader++, size1[i] = 1, size2[i] = 0;
else coder++, size1[i] = 0, size2[i] = 1;
}
ans = (reader * C(coder, 2) % p + C(coder, 3)) % p;
cout << ans << endl;
for(int i = 1; i <= n - 1; i++)
{
ll x, y, tmp = 0;
scanf("%lld%lld", &x, &y);
x = get(x), y = get(y); //一定记得这里是对树根进行操作！所以得先get
tmp = (tmp + size2[x] * size2[y] * (reader - size1[x] - size1[y])) % p;
tmp = (tmp + size2[x] * size2[y] * (coder - size2[x] - size2[y])) % p;
tmp = (tmp + size1[x] * size2[y] * (coder - size2[x] - size2[y])) % p;
tmp = (tmp + size2[x] * size1[y] * (coder - size2[x] - size2[y])) % p;
merge(x, y);
ans = (ans + p - tmp) % p;
cout << ans << endl;
}
}
}