51nod-1253: Kundu and Tree

【传送门：51nod-1253】

---

简要题意：

　　给出一棵n个点的树，树上的边要么为黑，要么为红

　　求出所有的三元组(a,b,c)的数量，满足a到b，b到c，c到a三条路径上分别有至少一条红边

---

题解：

　　显然黑边是没用的，那么我们将只有黑边相连的点分成若干的连通块

　　那么答案就很显然了，容斥一手

　　就是(所有三元组的数量)-(三个点都在一个连通块的数量)-(两个点在一个连通块，另一个不在的数量)

---

参考代码：

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<cmath>
using namespace std;
typedef long long LL;
LL Mod=1e9+7;
struct node
{
    int x,y,c,next;
}a[110000];int len,last[51000];
void ins(int x,int y,int c)
{
    a[++len]=(node){x,y,c,last[x]};
    last[x]=len;
}
int fa[51000];
int findfa(int x)
{
    if(fa[x]!=x) fa[x]=findfa(fa[x]);
    return fa[x];
}
int tot[51000];
void dfs(int x,int f,int rt)
{
    tot[rt]++;
    for(int k=last[x];k;k=a[k].next)
    {
        int y=a[k].y;
        if(y==f||a[k].c==0) continue;
        dfs(y,x,rt);fa[y]=rt;
    }
}
char st[5];
int main()
{
    int n;
    scanf("%d",&n);
    len=0;memset(last,0,sizeof(last));
    for(int i=1;i<n;i++)
    {
        int x,y;
        scanf("%d%d%s",&x,&y,st+1);
        ins(x,y,st[1]=='b');ins(y,x,st[1]=='b');
    }
    for(int i=1;i<=n;i++) fa[i]=i,tot[i]=0;
    for(int i=1;i<=n;i++) if(fa[i]==i) dfs(i,0,i);
    LL ans=(LL)n*(n-1)*(n-2)/6%Mod;
    for(int i=1;i<=n;i++)
    {
        int fx=findfa(i);
        if(fx==i)
        {
            if(tot[i]>=3) ans=(ans-(LL)tot[i]*(tot[i]-1)*(tot[i]-2)/6%Mod+Mod)%Mod;
            if(tot[i]>=2) ans=(ans-(LL)tot[i]*(tot[i]-1)/2%Mod*(n-tot[i])%Mod+Mod)%Mod;
        }
    }
    printf("%lld\n",ans);
    return 0;
}