codeforces 161D Distance in Tree 树形dp

题目链接：

http://codeforces.com/contest/161/problem/D

D. Distance in Tree

time limit per test 3 seconds
memory limit per test 512 megabytes
#### 问题描述
> A tree is a connected graph that doesn't contain any cycles.
>
> The distance between two vertices of a tree is the length (in edges) of the shortest path between these vertices.
>
> You are given a tree with n vertices and a positive number k. Find the number of distinct pairs of the vertices which have a distance of exactly k between them. Note that pairs (v, u) and (u, v) are considered to be the same pair.
#### 输入
> The first line contains two integers n and k (1 ≤ n ≤ 50000, 1 ≤ k ≤ 500) — the number of vertices and the required distance between the vertices.
>
> Next n - 1 lines describe the edges as "ai bi" (without the quotes) (1 ≤ ai, bi ≤ n, ai ≠ bi), where ai and bi are the vertices connected by the i-th edge. All given edges are different.
#### 输出
> Print a single integer — the number of distinct pairs of the tree's vertices which have a distance of exactly k between them.
>
> Please do not use the %lld specifier to read or write 64-bit integers in С++. It is preferred to use the cin, cout streams or the %I64d specifier.
#### 样例
> **sample input**
> 5 2
> 1 2
> 2 3
> 3 4
> 2 5
>
> **sample output**
> 4

题意

给你一颗树，每条边长为1，求所有距离为k的顶点对，(u,v)和(v,u）算一对。

题解

树形dp：

dp[i][j]表示与第i个节点距离为j的节点数。

两次dfs:

第一次求以i为根的子树中与i距离为j的节点数dp[i][j]。

第二次求i与不在i的子树中的节点金额距离为j的节点数。

两次加起来就是表示与i节点距离为j的所有的树上节点数。

答案就是sigma(dp[i][k])。

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<queue>
#include<map>
#define lson (o<<1)
#define rson ((o<<1)|1)
#define M (l+(r-l)/2)
using namespace std;

const int maxn=5e4+10;
const int maxm=555;

typedef __int64 LL;

int n,k;
LL dp[maxn][maxm];
vector<int> G[maxn];

void dfs1(int u,int fa) {
	dp[u][0]=1;
	for(int i=0;i<G[u].size();i++){
		int v=G[u][i];
		if(v==fa) continue;
		dfs1(v,u);
		for(int j=0;j+1<=k;j++){
			dp[u][j+1]+=dp[v][j];
		}
	}
}

LL tmp[maxm];
void dfs2(int u,int fa) {
	if(fa!=-1){
		tmp[0]=dp[fa][0];
		for(int j=1;j<=k;j++){
			tmp[j]=dp[fa][j]-dp[u][j-1];
		}
		for(int j=0;j+1<=k;j++){
			dp[u][j+1]+=tmp[j];
		}
	}
	for(int i=0;i<G[u].size();i++){
		int v=G[u][i];
		if(v==fa) continue;
		dfs2(v,u);
	}
}

int main() {
	scanf("%d%d",&n,&k);
	memset(dp,0,sizeof(dp));
	for(int i=0; i<n-1; i++) {
		int u,v;
		scanf("%d%d",&u,&v);
		G[u].push_back(v);
		G[v].push_back(u);
	}
	dfs1(1,-1);
	dfs2(1,-1);
	LL ans=0;
	for(int i=1;i<=n;i++){
		ans+=dp[i][k];
	}
	printf("%I64d\n",ans/2);
	return 0;
}