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