题解 CF109C 【Lucky Tree】

思路：

\(\quad\)树形\(DP +\) 容斥 ， \(f[x]\) 表示以 \(x\) 为根的子树中有几个点到 \(x\) 的路径包含幸运边， \(g[x]\) 表示除了 \(x\) 的子树外有几个点到 \(x\) 的路径包含幸运边，最后统计答案就是 \(ans=\sum_{i=1}^n\) \((f[i]\times (f[i]-1)+g[i]\times (g[i]-1)+2\times f[i]\times g[i])\)，数组 \(f\) 和 \(g\) 的转移也很简单，当 \(y\) 是 \(x\)的儿子时，若 \(x\) 和 \(y\) 的连边为幸运边，那么 \(f[x]+=size[y]\) ， \(y\) 的子树(包括 \(y\))到 \(x\) 的路径都包含幸运边，否则 \(f[x]+=f[y]\) ， \(g\) 的转移也类似，如果x和y的连边是幸运边，那么 \(g[y]=(n-size[y])\) ，否则 \(g[y]=g[x]+f[x]-f[y]\) ，两遍 \(dfs\) 即可完成的操作，一遍自下而上，一遍自上而下，时间复杂度为 \(O(3n)\) 。

#include<iostream>
#include<cstdio>
#include<string>
#include<algorithm>
#include<cstring>
using namespace std;
#define re register int
#define int long long
#define LL long long
#define il inline
#define next nee
#define inf 1e18
il int read()
{
  int x=0,f=1;char ch=getchar();
  while(!isdigit(ch)&&ch!='-')ch=getchar();
  if(ch=='-')f=-1,ch=getchar();
  while(isdigit(ch))x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
  return x*f;
}
il void print(int x)
{
  if(x<0)putchar('-'),x=-x;
  if(x/10)print(x/10);
  putchar(x%10+'0');
}
const int N=1e5+5;
int n,m,next[N<<1],go[N<<1],head[N],tot,father[N],size[N],f[N],g[N],ans;
bool s[N<<1];
il bool check(int x)//判断是否是幸运数字
{
  while(x)
    {
      if(x%10!=7&&x%10!=4)return 0;
      x/=10;
    }
  return 1;
}
il void Add(int x,int y,bool z)
{next[++tot]=head[x];head[x]=tot;go[tot]=y;s[tot]=z;}
il void dfs1(int x,int fa,bool z)//第一遍dfs，处理f数组
{
  father[x]=fa;size[x]=1;
  for(re i=head[x],y;i,y=go[i];i=next[i])
    {
      if(y==fa)continue;
      dfs1(y,x,s[i]);
      size[x]+=size[y];
      if(s[i])f[x]+=size[y];
      else f[x]+=f[y];
    }
}
il void dfs2(int x)//第二遍dfs，处理g数组
{
  for(re i=head[x],y;i,y=go[i];i=next[i])
    {
      if(y==father[x])continue;
      if(s[i])g[y]=n-size[y];
      else g[y]=g[x]+f[x]-f[y];
      dfs2(y);
    }
}
signed main()
{
  n=read();
  for(re i=1;i<n;i++)
    {re x=read(),y=read(),z=check(read());Add(x,y,z);Add(y,x,z);}
  dfs1(1,0,0);dfs2(1);
  for(re i=1;i<=n;i++)
    ans+=f[i]*(f[i]-1)+g[i]*(g[i]-1)+f[i]*g[i]*2;//简单的组合数学计算方案
  print(ans);
  return 0;
}