109 C Lucky Tree

問題

n個の頂点からなる重み付き無向木が与えられる。
lucky edgeとは、辺の重みの数字を10進法で書いたときすべて4か7である辺を言う。


このとき、次の三つ組の個数を求めよ。
(i,j,k) i->jへ、lucky edgeを通るパスがある かつ i -> kへlucky edgeを通るパスがある。

制約条件

n≦10^5, w≦10^9

方針

ツリーDP.
今頂点iにいるとする。
自分の子ノード以下でlucky edgeを通って行くノードの数をdp[i]とする。
自分の親ノードを経由してlukcy edgeを通って行くノードの数をdp2[i]とする。


すると、dp[i]は、Σ(子ノードj) i-jがluckyならjの子の数, そうでないならdp[j]
dp2[j]はi-jがluckyなら全体-(j以下の子ノード)、そうでないならdp2[j]+Σ(j以外の子[k])dp[k]となる。


dp,dp2を求め終えたら、求める数は、
Σdp[i]*(dp[i]-1)+(dp2[i]*dp2[i]-1)+2*dp[i]*dp2[i]となる。
(iをひとつ決めて、子側からj,kを選ぶ場合、親側からj,kを選ぶ場合、子と親側からj,kを選ぶ場合を足す、これを全てのiについて足し合わせる)

ソースコード

int dp[100000],dp2[100000];
int n, cld[100000];
vector<vector<pi> >e;

bool islucky(int n){
  for(;n;n/=10)if(n%10!=4&&n%10!=7)return 0;
  return 1;
}
void _(int c,int p){
  cld[c]=1;
  fr(i,e[c])if(i->first!=p){
    _(i->first,c);
    cld[c]+=cld[i->first];
  }
}
void rec(int c,int p){
  fr(i,e[c])if(i->first!=p){
    rec(i->first,c);
    if(i->second)dp[c]+=cld[i->first];
    else dp[c]+=dp[i->first];
  }
}
void rec2(int c,int p){
  fr(i,e[c])if(i->first!=p){
    if(i->second)dp2[i->first]=cld[0]-cld[i->first];
    else dp2[i->first]=dp2[c]+dp[c]-dp[i->first];
    rec2(i->first,c);
  }
}

void run(){
  cin>>n;
  e.resize(n);
  
  rep(i,n-1){
    int a,b,c; cin>>a>>b>>c;
    a--; b--; c=islucky(c);
    e[a].pb(mp(b,c)); e[b].pb(mp(a,c));
  }
  _(0,-1); rec(0,-1); rec2(0,-1);
  
  ll ans=0;
  rep(i,n)ans+=dp[i]*(dp[i]-1ll)+dp2[i]*(dp2[i]-1ll)+2ll*dp[i]*dp2[i];
  cout<<ans<<endl;
}