原文链接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; }