51Nod1253 Kundu and Tree 容斥原理

​​题目传送门 - 51Nod1253​​题意

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

题解

　　把黑色边连接的点搞成一块。

　　答案 = 任选 3 个点的方案数 - 在同一个黑色块中选 3 个点的方案数 - 任选三个数，其中两个点在同一个黑色块中的方案数。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N=50005;
int read(){
  int x=0;
  char ch=getchar();
  while (!isdigit(ch))
    ch=getchar();
  while (isdigit(ch))
    x=(x<<1)+(x<<3)+ch-48,ch=getchar();
  return x;
}
struct Gragh{
  static const int M=N*2;
  int cnt,y[M],z[M],nxt[M],fst[N];
  void clear(){
    cnt=0;
    memset(fst,0,sizeof fst);
  }
  void add(int a,int b,int c){
    y[++cnt]=b,z[cnt]=c,nxt[cnt]=fst[a],fst[a]=cnt;
  }
}g;
int n,dsize[N];
vector <int> sz;
void dfs(int x,int pre){
  dsize[x]=1;
  for (int i=g.fst[x];i;i=g.nxt[i])
    if (g.y[i]!=pre){
      int y=g.y[i];
      dfs(y,x);
      if (g.z[i])
        dsize[x]+=dsize[y];
      else
        sz.push_back(dsize[y]);
    }
}
LL calc(LL x){
  return x*(x-1)*(x-2)/6;
}
int main(){
  n=read();
  for (int i=1;i<n;i++){
    int x=read(),y=read();
    char s[2];
    scanf("%s",s);
    g.add(x,y,s[0]=='b');
    g.add(y,x,s[0]=='b');
  }
  sz.clear();
  dfs(1,0);
  sz.push_back(dsize[1]);
  LL ans=calc(n);
  for (int i=0;i<sz.size();i++)
    ans-=calc(sz[i])+1LL*sz[i]*(sz[i]-1)/2*(n-sz[i]);
  printf("%lld",ans%1000000007);
  return 0;
}