【51nod 1253】Kundu and Tree

Description

树包含N个点和N-1条边。树的边有2中颜色红色(‘r’)和黑色（’b’)。给出这N-1条边的颜色，求有多少节点的三元组(a,b,c)满足：节点a到节点b、节点b到节点c、节点c到节点a的路径上，每条路径都至少有一条边是红色的。
注意(a,b,c), (b,a,c)以及所有其他排列被认为是相同的三元组。输出结果对1000000007取余的结果。

Solution

这道题看一下就觉得是容斥。现在我们考虑所有黑边都是没用的，所以我们用并查集构出一个个全是黑边的连通块。最后答案ans=任意选三个点的方案-在一个连通块中选三个点-在一个连通块中选两个在另一连通块选一个点的方案。

Code

#include<iostream>
#include<cmath>
#include<cstring>
#include<cstdio>
#include<algorithm>
#define ll long long
using namespace std;
const ll mo=1e9+7;
const int maxn=1e5+5; 
ll fa[maxn],size[maxn],a[maxn];
ll n,i,t,j,k,l,x,y,z,num,sum,ans;
char s[10];
int getfather(int x){
    if (x==fa[x])return x;
    fa[x]=getfather(fa[x]);return fa[x];
}
ll make(ll x){
    return x*(x-1)*(x-2)/6;
}
int main(){
//  freopen("data.in","r",stdin);freopen("data.out","w",stdout);
    scanf("%lld",&n);
    for (i=1;i<=n;i++)fa[i]=i,size[i]=1;
    for (i=1;i<n;i++){
        scanf("%lld%lld%s\n",&x,&y,s+1);
        if (s[1]=='b'){
            x=getfather(x);y=getfather(y);
            if (x!=y) size[x]+=size[y],fa[y]=x;
        }
    }
    for (i=1;i<=n;i++)
        if (fa[i]==i) a[++num]=size[i],sum+=size[i],ans-=make(size[i]);
    ans+=make(sum);
    for (i=1;i<=num;i++)
        ans-=a[i]*(a[i]-1)/2*(sum-a[i]);
    ans%=mo;
    printf("%lld\n",ans);
}