CodeForces 109C 树形DP Lucky Tree

赶脚官方题解写得挺清楚的说，=_=

注意数据范围用long long，否则会溢出。

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

const int maxn = 100000 + 10;

int n;
vector<int> G[maxn], C[maxn];

inline int lucky(int x)
{
    while(x)
    {
        if(x % 10 != 4 && x % 10 != 7) return 0;
        x /= 10;
    }
    return 1;
}

long long d[maxn], sz[maxn], f[maxn];

void dfs(int u, int fa)
{
    sz[u] = 1;
    f[u] = 0;
    for(int i = 0; i < G[u].size(); i++)
    {
        int v = G[u][i];
        if(v == fa) continue;
        dfs(v, u);
        sz[u] += sz[v];
        if(C[u][i]) f[u] += sz[v];
        else f[u] += f[v];
    }
}

void dfs2(int u, int fa)
{
    for(int i = 0; i < G[u].size(); i++)
    {
        int v = G[u][i];
        if(v == fa) continue;
        if(C[u][i]) d[v] = sz[1] - sz[v];
        else d[v] = d[u] + f[u] - f[v];
        dfs2(v, u);
    }
}

int main()
{
    //freopen("in.txt", "r", stdin);

    scanf("%d", &n);
    for(int i = 1; i < n; i++)
    {
        int u, v, d; scanf("%d%d%d", &u, &v, &d);
        int t = lucky(d);
        G[u].push_back(v); C[u].push_back(t);
        G[v].push_back(u); C[v].push_back(t);
    }

    dfs(1, 0);
    dfs2(1, 0);

    long long ans = 0;
    for(int i = 1; i <= n; i++) ans += f[i] * (f[i] - 1) + d[i] * (d[i] - 1) + 2 * f[i] * d[i];
    printf("%I64d
", ans);

    return 0;
}