Codeforces 161 D. Distance in Tree (树dp)

Posted 2020-08-13 Recoder

题目链接：http://codeforces.com/problemset/problem/161/D

题意：

        给你一棵树，问你有多少对点的距离为k。

思路：

        dp[i][j]表示离i节点距离为j的点个数，2次dfs，一次从底向上，另一次从顶向下。

 1 //#pragma comment(linker, "/STACK:102400000, 102400000")
 2 #include <algorithm>
 3 #include <iostream>
 4 #include <cstdlib>
 5 #include <cstring>
 6 #include <cstdio>
 7 #include <vector>
 8 #include <cmath>
 9 #include <ctime>
10 #include <list>
11 #include <set>
12 #include <map>
13 using namespace std;
14 typedef long long LL;
15 typedef pair <int, int> P;
16 const int N = 1e5 + 5;
17 int dp[50005][505], n, m;
18 vector <int> G[50005];
19 
20 void dfs1(int u, int p) {
21     dp[u][0] = 1;
22     for(int i = 0; i < G[u].size(); ++i) {
23         int v = G[u][i];
24         if(v == p) 
25             continue;
26         dfs1(v, u);
27         for(int i = 1; i <= m; ++i) {
28             dp[u][i] += dp[v][i - 1];
29         }
30     }
31 }
32 
33 void dfs2(int u, int p) {
34     for(int i = 0; i < G[u].size(); ++i) {
35         int v = G[u][i];
36         if(v == p)
37             continue;
38         for(int i = m; i >= 2; --i) {
39             dp[v][i] += dp[u][i - 1] - dp[v][i - 2];
40         }
41         dp[v][1]++;
42         dfs2(v, u);
43     }
44 }
45 
46 int main()
47 {
48     int u, v;
49     while(~scanf("%d %d", &n, &m)) {
50         for(int i = 1; i < n; ++i) {
51             scanf("%d %d", &u, &v);
52             G[u].push_back(v);
53             G[v].push_back(u);
54         }
55         if(m == 1) {
56             printf("%d\n", n - 1);
57         } else {
58             dfs1(1, -1);
59             dfs2(1, -1);
60             LL ans = 0;
61             for(int i = 1; i <= n; ++i) {
62                 ans += (LL)dp[i][m];
63             }
64             printf("%lld\n", ans / 2);
65         }
66     }
67     return 0;
68 }