题目描述
懒惰的温温今天上班也在偷懒。盯着窗外发呆的温温发现，透过窗户正巧能看到一棵n个节点的树。一棵n个节点的树包含n-1条边，且n个节点是联通的。树上两点之间的距离即两点之间的最短路径包含的边数。
温温瞧见树上有一只灵活的小松鼠正在节点间跳来跳去。观察了很长一段时间之后，温温发现：松鼠每次跳跃，从当前节点可以跳到任何与当前节点距离不超过k的节点。
突发奇想的温温想要知道，如果定义f(u, v)为从松鼠从u点到v点所需的最少跳跃次数，那么，对于树上的所有点对(u, v)，f(u, v)的总和是多少。
注意：(u, v)和(v, u)视作同一个点对，只计算一次答案。
输入
第一行两个整数n和k。
接下来n-1行每行两个整数ai, bi，表示节点ai和bi之间存在一条边。
1 ≤ k ≤ 5
2 ≤ n ≤ 5000 for 40%
2 ≤ n ≤ 200000 for 100%
输出
输出一个整数，表示所求的f(u, v)总和。
样例输入 Copy
【样例1】
6 2
1 2
1 3
2 4
2 5
4 6
【样例2】
13 3
1 2
3 2
4 2
5 2
3 6
10 6
6 7
6 13
5 8
5 9
9 11
11 12
样例输出 Copy
【样例1】
20
【样例2】
114
#include <bits/stdc++.h>
using namespace std ;
typedef long long ll ;
const int N = 2e5 + 10 ;
int e[N << 1] , ne[N << 1] , h[N << 1] , idx , flow[N] , in[N] , dp[N][6] , size[N] , k , b[6] , n ;
void add(int a, int b)
{
e[idx] = b , ne[idx] = h[a] , h[a] = idx ++ ;
}
void dfs1(int u , int fa)
{
memset(b , 0 , sizeof b) ;
if(fa)
{
b[0] = dp[fa][0] - dp[u][k - 1] - size[u] ;
for(int i = 1; i < k ;i ++) b[i] = dp[fa][i] - dp[u][i - 1] ;
}
dp[u][0] += b[k - 1] + n - size[u] ;
for(int i = 1 ; i < k ;i ++) dp[u][i] += b[i - 1] ;
for(int i = h[u] ; ~i ; i = ne[i])
{
int v = e[i] ;
if(v == fa) continue ;
dfs1(v , u) ;
}
}
void dfs(int u , int fa)
{
size[u] = 1 ;
for(int i = h[u] ; ~i ; i = ne[i])
{
int v = e[i] ;
if(v == fa) continue ;
dfs(v , u) ;
size[u] += size[v] ;
dp[u][0] += dp[v][k - 1] + size[v] ;
for(int i = 1; i < k ;i ++) dp[u][i] += dp[v][i - 1] ;
}
return ;
}
void work()
{
cin >> n >> k ;
memset(flow , 0 , sizeof flow) ;
idx = 0 ;
memset(h , -1 , sizeof h) ;
for(int i = 1 , a , b , c ; i < n ;i ++)
cin >> a >> b ,
add(a , b) , add(b , a);
dfs(1, 0) ;
dfs1(1 , 0) ;
ll res = 0 ;
for(int i = 1; i <= n ;i ++) res += 1ll * dp[i][0] ;
cout << res / 2 << endl ;
return ;
}
int main()
{
work() ;
return 0 ;
}