51Nod1253 Kundu and Tree 容斥原理

原文链接https://www.cnblogs.com/zhouzhendong/p/51Nod1253.html

题目传送门 - 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;
}