POI 2014 HOTELS (树形DP)

题目链接 HOTELS

依次枚举每个点，以该点为中心扩展。

每次枚举的时候，从该点的儿子依次出发，搜完一个儿子所有的点之后进行答案统计。

这里用了一个小trick。

#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b)              for(int i(a); i <= (b); ++i)
#define for_edge(i, x)            for(int i = H[x]; i; i = X[i])

int n, x, y, et = 0;
int E[100010], X[100010], H[100010], f[5010], g[5010], tmp[5010];
long long ans;

inline void addedge(int a, int b){
    E[++et] = b, X[et] = H[a], H[a] = et;
    E[++et] = a, X[et] = H[b], H[b] = et;
}

void dfs(int x, int fa, int dep){
    ++tmp[dep];
    for_edge(i, x){
        int u = E[i];
        if (u != fa) dfs(u, x, dep + 1);
    }
}

int main(){

    scanf("%d", &n);
    rep(i, 1, n - 1){
        scanf("%d%d", &x, &y);
        addedge(x, y);
    }

    rep(u, 1, n){
        memset(f, 0, sizeof f);
        memset(g, 0, sizeof g);
        for_edge(i, u){
            memset(tmp, 0, sizeof tmp);
            x = E[i];
            dfs(x, u, 1);
            rep(j, 1, n){
                ans += (long long)g[j] * tmp[j];
                g[j] += f[j] * tmp[j];
                f[j] += tmp[j];
            }
        }
    }

    printf("%lld\n", ans);
    return 0;

}