1. 程式人生 > 實用技巧 >題解 CF109C 【Lucky Tree】

題解 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;
}