[Hackerrank题目选做] Kundu and Tree 组合数学

题目链接: https://www.hackerrank.com/challenges/kundu-and-tree

题目大意: 给一棵n个点的树，每条边被染成黑色或红色，

                  求三元组(x,y,z)的对数，使得(x,y),(y,z),(x,z)最短路径中都至少有一条红色的边。

题解: 考虑反向求解，答案=总的三元组对数-使得(x,y)(y,z)(x,z)最短路径中至少有一条全黑的三元组(x,y,z)对数.

          总的三元组对数：n*(n-1)*(n-2)/6

          使得(x,y)(y,z)(x,z)最短路径中至少有一条全黑的三元组(x,y,z)对数的求解，可以用树形dp+组合数学解决.

          先O(n)做一次树形dp，将树拆分为只有黑色边组成的森林，

          这样，任意连通的两点最短路径上都没有红色的边，而任意不连通的两点最短路径都必须经过红边。

          由此，每个不符合条件的三元组(x,y,z)满足，至少有两点在一个连通块中。

          这样，对于每个大小为t的联通块，s+=C(t,3)+C(t,2)*(n-t).

          输出n*(n-1)*(n-2)/6-s即为最终的答案。

Code:

#include <bits/stdc++.h>
#define ll long long
#define mod 1000000007ll
using namespace std;
int head[100005],nxt[200005],v[200005],w[200005],tot=0;
int sz[100005];bool g[100005];
inline void add(int a,int b,int val)
{tot++;nxt[tot]=head[a];head[a]=tot;v[tot]=b;w[tot]=val;}
void dfs(int pos,int fa)
{sz[pos]=1;
for (int i=head[pos];i;i=nxt[i])
{if (v[i]!=fa)
{dfs(v[i],pos);
if (w[i]) {g[v[i]]=1;}
else {g[v[i]]=0;sz[pos]+=sz[v[i]];}
}
}
}
int main (){
	int u,v,n,i;
	char s[5];
	scanf ("%d",&n);
	for (i=1;i<n;i++)
	{scanf ("%d%d",&u,&v);
	scanf ("%s",s);
	if (s[0]=='r') {add(u,v,1);add(v,u,1);}
	else {add(u,v,0);add(v,u,0);}
	}
	g[1]=1;
	dfs(1,0);
	ll ans=((ll)(n))*((ll)(n-1))*((ll)(n-2));
	ans/=6ll;
	ll si=0;
	for (i=1;i<=n;i++)
	{if (g[i]&&sz[i]>=2)
	{si+=(((ll)(sz[i]))*((ll)(sz[i]-1))*((ll)(sz[i]-2)))/6ll;
	si+=(((ll)(sz[i]))*((ll)(sz[i]-1))/2ll)*((ll)(n-sz[i]));
	}
	}
	ans-=si;
	ans%=mod;
	printf ("%lld\n",ans);
	return 0;
}