BZOJ 3522 POI 2014 Hotel 树形DP

题目大意

给出一棵树，问选择三个点，使得这三个点相互的距离相等的方案有多少种。

思路

这三个点肯定不能再一条链上， 那么就肯定能够确定一个中心点，使得三个点到这个中心点的距离都相等。
之后我们就可以枚举这个中心点，对于每个深度统计一下就可以了。虽然看起来像是 O(n3) 的，但是跑的飞起啊。

CODE

#define _CRT_SECURE_NO_WARNINGS

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

int points;
int head[MAX], total;
int _next[MAX << 1], aim[MAX << 1];

inline void Add(int x, int y)
{
    _next[++total] = head[x];
    aim[total] = y;
    head[x] = total;
}

int deep[MAX];
int now;

void DFS(int x, int last, int d)
{
    ++deep[d];
    for(int i = head[x]; i; i = _next[i]) {
        if(aim[i] == last)  continue;
        DFS(aim[i], x, d + 1);
    }
}

long long cnt[MAX], ans[MAX];

int main()
{
    cin >> points;
    for(int x, y, i = 1; i < points; ++i) {
        scanf("%d%d", &x, &y);
        Add(x, y), Add(y, x);
    }
    long long output = 0;
    for(int x = 1; x <= points; ++x) {
        memset(cnt, 0, sizeof(cnt));
        memset(ans, 0, sizeof(ans));
        for(int i = head[x]; i; i = _next[i]) { 
            memset(deep, 0, sizeof(deep));
            DFS(aim[i], x, 1);
            for(int d = 1; d <= points; ++d) {
                output += ans[d] * deep[d];
                ans[d] += cnt[d] * deep[d];
                cnt[d] += deep[d];
            }
        }
    }
    cout << output << endl;
    return 0;
}