CF 161D Distance in Tree 樹形DP

題意：給出一個無根樹，統計有多少節點對之前的距離恰好等於K。相鄰節點間的距離爲1.
思路：樹形DP。
我們可以將無根樹轉化成有根樹，這個很容易做到。
對於一個子樹，其根結點爲s，我們只需統計s的子樹之間距離爲K節點對即可。然後對每個子樹都進行相同的統計，就可以得到最終的答案。
因爲這些節點對是經過根結點s，所以他們和根結點的距離之和就是K。所以我們要知道在以s爲根的子樹下，有多少個距離爲i的節點。然後在利用計數的乘法原理就可以了。
但是還要考慮的一個問題是如何去重。只是簡單的統計出節點s的所有兒子的距離分佈，會讓我們喪失太多的信息，因爲我們不知道這些兒子位於哪棵子樹中。因爲求節點s的所有兒子的距離的時候，是按照一個一個子樹來求的，這個過程已經非常自然的幫我們把子樹劃分開。所以，我們要在這個過程中先求出當前子樹和前面所有子樹產生了多少滿足條件的節點對，再把當前子樹的節點加入到根結點上。

#include <cstdio>
#include <algorithm>
#include <cstring>

using namespace std;

const int MAX = 50010;

int N,K;
int to[MAX*2],nxt[MAX*2],head[MAX],tot;
long long s[50010][510];

void init()
{
    memset(head,-1,sizeof(head));
    tot = 0;
}
void addedge(int u, int v)
{
    to[tot] = v, nxt[tot] = head[u];
    head[u] = tot++;
    to[tot] = u, nxt[tot] = head[v];
    head[v] = tot++;
}

long long dfs(int u, int p)
{
    s[u][0] = 1;
    long long ans = 0;
    for(int i = head[u]; ~i ; i = nxt[i]){
        int v = to[i];
        if(p == v) continue;
        ans += dfs(v,u);
        for(int i = 0; i < K; ++i)
            ans += s[u][i] * s[v][K - i - 1];
        for(int i = 1; i < K; ++i)
            s[u][i] += s[v][i-1];
    }
    return ans;
}


int main(void)
{
    //freopen("input.txt","r",stdin);
    init();
    scanf("%d%d",&N,&K);
    for(int i = 0; i < N - 1; ++i){
        int u, v;
        scanf("%d%d",&u,&v);
        addedge(u,v);
    }
    printf("%I64d\n",dfs(1,0));
    return 0;
}