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

include

include

include

include

include

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 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;
}