【51Nod1588】幸运树（树形DP）

题目描述：

幸运树传送门

算法：

树形DP 求方案数

做法：

分类

要统计方案数，首先应给这些点分类，才能计数。怎么分类既方便有不会有重复呢？这样分：

当算 u 这个点时，令 f[u] 为在 u 的子树内，到 u 的路径中有幸运边的点有几个，令 g[u] 为在 u 的子树外，到 u 的路径中有幸运边的点有几个。这样当我们计算有 ( i , j , k ) 的个数时，我们枚举 i ，当 i=u 时：

$$( u , j , k ) 的方案数 = f[u]*(f[u]-1)+g[u]*(g[u]-1)+f[u]*g[u]*2$$

转移

这样分类很好，实际上也很好计算 f 和 g，因为这是树形DP，只要与当前子树有关，就好转移（我是这么想的，也许是因为我做题还不够）。假设 v 是 u 的一个孩子，fa 是 u 的父亲：

$$边( u , v ) 是幸运边：f[u] += size[v]$$

$$边( u , v ) 不是幸运边：f[u]+=f[v]$$

$$边( fa, u ) 是幸运边：g[u] = size[1]-size[u] ( 1是根 )$$

$$边( fa, u ) 不是幸运边：g[u] = g[fa]+f[fa]-f[u]$$

这么转移实在是强，也许这是一个套路，也许我还没掌握住。

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;

const int N=100010;
int n, mi;
int head[N];

struct Edge{
    int to, next, w;
}e[N<<1];

inline void add(int u,int v,int w){
    Edge E = {v,head[u],w};
    e[++mi] = E;
    head[u] = mi;
}

int l, r;
char s[15];
int q[N], fa[N], size[N], f[N], g[N];

inline void bfs(){          // Eirlys 大佬用 bfs ，我也跟着学了学 。。。 
    l = r = 1;
    q[r++] = 1;
    int u, v, p;
    while(l<r){
        u = q[l]; size[u]=1; l++;
        for(p=head[u]; p; p=e[p].next) if(e[p].to!=fa[u]){
            v = e[p].to;
            fa[v] = u;      // 计算 fa[] 
            q[r++] = v;
        }
    }
    for(int i=n; i>=1; --i){
        u = q[i];
        for(p=head[u]; p; p=e[p].next) if(e[p].to!=fa[u]){
            v = e[p].to;
            size[u] += size[v];                                     // 计算 size[] 
            f[u] += e[p].w ? size[v] : f[v];                        // 计算 f[] 
        }
    }
    for(int i=1; i<=n; ++i){
        u = q[i];
        for(p=head[u]; p; p=e[p].next) if(e[p].to!=fa[u]){
            v = e[p].to;
            g[v] = e[p].w ? size[1]-size[v] : g[u]+f[u]-f[v];       // 计算 g[] 
        }
    }
}

int main(){
    scanf("%d",&n);
    for(int i=1, u, v, w, len, j; i<n; ++i){
        scanf("%d%d%s",&u, &v, s);
        len = strlen(s);
        w=1;
        for(j=0; j<len; ++j) if(s[j]!='4'&&s[j]!='7'){ w=0; break; }
        add(u,v,w); add(v,u,w);
    }
    bfs();
    long long ans = 0;
    for(int i=1; i<=n; ++i) ans += f[i]*1ll*(f[i]-1)+g[i]*1ll*(g[i]-1)+f[i]*1ll*g[i]*2;     // 统计答案  
    printf("%lld\n",ans);
    while(1);
}