dp专题 51Nod - 1588

题目
比得喜欢幸运数字。这里所说的幸运数字是由4和7组成的正整数。比如，数字47，744，4是幸运数字，而5，17，467就不是。

一天，比得遇到一棵由n个点组成的树。另外，这棵树是带权的，即每条边有一个权值（由一个正整数表示）。如果一条边的权值是一个幸运数字，那么我们就说这条边是一条幸运边。说明一下，一棵n个结点的树是由n个结点和n-1条边组的无环的无向图。

比得好奇，在树中有多少个满足以下条件的三元组tr（i,j,k）（i,j,k是三个不同的点）。

1.i到j有路径，i到k也有路径

2.每条路径中至少有一条幸运边。

数字的顺序是有意义的，举例说明，tr(1,2,3)，tr(1,3,2),tr(2,1,3)是三个不同的序列。

现在要求计算在树中存在多少个这样的三元组关系。

样例解释：

样例一中的16种情况分别为:

(1,2,4),(1,4,2),(2,1,3),(2,1,4),(2,3,1),(2,3,4),(2,4,1),(2,4,3),(3,2,4),(3,4,2),(4,1,2),(4,1,3),(4,2,1),(4,2,3),(4,3,1),(4,3,2)

Input
单组测试数据 第一行包含一个整数n（1≤n≤10^5）。 接下来的n-1行中每行有三个整数 ui vi wi (1≤ui,vi≤n,1≤wi≤10^9) 分别表示有边相连的两个点和这条边的权值。
Output
共一行，表示题目中所要计算的三元组的个数。
Sample Input
4
1 2 4
3 1 2
1 4 7
Sample Output
16
题目大意
这是一棵树,不妨设根节点为1.
题目要求找幸运数,对于节点x,如果能算出有多少个点到x至少有一条幸运边,那么从这些点取两个点进行排列的数量即为以点x为中间数的幸运数数量。
(由于这是一棵树,可以先算出x的子树到点x有至少一条幸运边的数量加上除x的子树外其他节点到x节点至少一条幸运边的数量)。

进行一次dfs遍历,可以算出x点的子树有多少个节点到x点至少一条幸运边(设为ss[x]),多少个节点到x点没有幸运边(fa[x])
则对于x的孩子son,

 如果x到son有幸运边,则ss[x]+=ss[son]+fa[son]+1（加上son到x的那一条）
 如果x到son没有幸运边,则ss[x]+=ss[son],fa[x]+=fa[son]+1(加上son到x的那一条)

现在对于节点x,以及它的一个孩子节点son,设wss[x]为除其子树外其他节点到x节点至少一条幸运边的数量,wfa[x]为除其子树外其他节点到x节点不存在幸运边的数量。(1为根节点,wss[1]=0,wfa[1]=0)

 如果x到son为幸运边
 x的除son及其子节点外的其他子节点到x有至少一条幸运边的数量为ss[x]-(ss[son]+fa[son]+1)
 x的除son及其子节点外的其他子节点到x没有幸运边的数量为fa[x]
 wss[x]=ss[x]+fa[x]-(ss[son]+fa[son+1)+1
 wfa[x]=0
 如果x到son不是幸运边
  x的除son及其子节点外的其他子节点到x有至少一条幸运边的数量为ss[x]-son[x]
  x的除son及其子节点外的其他子节点到x没有幸运边的数量为fa[x]-(fa[son]+1)
  wss[x]=ss[x]-son[x]
  faa[x]=fa[x]-(fa[son]+1)+1

代码

#include <cstdio>
#include <vector>
#include <iostream>
using namespace std;
int n;
long long ans;
long long ss[101010],fa[101010],wss[101010],wfa[101010];
vector <long long> a[101010],b[101010];
bool check(long long x)
{
	while (x!=0)
	{
		if (x%10!=4 && x%10!=7) return false;
		x/=10;
	}
	return true;
} 
void dfs1(int x,int father)
{
	for (int i=0;i<a[x].size();i++)
	{
		int son=a[x][i];
		if (son!=father)
		{
			dfs1(son,x);
			if (b[x][i]) ss[x]+=ss[son]+fa[son]+1;
			else 
			{
				ss[x]+=ss[son];
				fa[x]+=fa[son]+1;
			}
			
		}
	}
}
void dfs2(int x,int father,long long sss,long long faa,int symt)
{
	if (symt==0) 
	{
		wss[x]=sss;
		wfa[x]=faa+1;
	}
    if (symt==1) 
	{
	    wss[x]=sss+faa+1;
	}
	if (symt!=3)
	{
	sss=wss[x]+ss[x];
	faa=wfa[x]+fa[x];
    }
	for (int i=0;i<a[x].size();i++)
	  {
	  	int son=a[x][i];
	  	if (son!=father)
	  	{
	  		if (b[x][i]) dfs2(son,x,sss-(fa[son]+ss[son]+1),faa,1);
	  		else dfs2(son,x,sss-ss[son],faa-(fa[son]+1),0);
	  	}
	  }
}
int main()
{
	freopen("n.in","r",stdin);
	freopen("n.out","w",stdout);
	scanf("%d",&n);
	for (int i=1;i<=n-1;i++)
	{
		long long u,v,w;
		cin>>u>>v>>w;
		bool x=check(w);
		a[u].push_back(v);
		b[u].push_back(x);
		a[v].push_back(u);
		b[v].push_back(x);
	}
    dfs1(1,0);
    dfs2(1,0,ss[1],fa[1],3);
    for (int i=1;i<=n;i++)
      ans+=ss[i]*(ss[i]-1)+wss[i]*(wss[i]-1)+ss[i]*wss[i]*2;
	cout<<ans<<endl;
}