【CF771C】Bear and Tree Jumps

题目

题目描述
A tree is an undirected connected graph without cycles. The distance between two vertices is the number of edges in a simple path between them.

Limak is a little polar bear. He lives in a tree that consists of nn vertices, numbered 11 through nn .

Limak recently learned how to jump. He can jump from a vertex to any vertex within distance at most kk .

For a pair of vertices (s,t)(s,t) we define f(s,t)f(s,t) as the minimum number of jumps Limak needs to get from ss to tt . Your task is to find the sum of f(s,t)f(s,t) over all pairs of vertices (s,t)(s,t) such that s<t .

输入格式
The first line of the input contains two integers nn and kk ( 2<=n<=2000002<=n<=200000 , 1<=k<=51<=k<=5 ) — the number of vertices in the tree and the maximum allowed jump distance respectively.

The next n-1n−1 lines describe edges in the tree. The ii -th of those lines contains two integers a_{i}a
i
​
and b_{i}b
i
​
( 1<=a_{i},b_{i}<=n1<=a
i
​
,b
i
​
<=n ) — the indices on vertices connected with ii -th edge.

It’s guaranteed that the given edges form a tree.

输出格式
Print one integer, denoting the sum of f(s,t)f(s,t) over all pairs of vertices (s,t)(s,t) such that s<t .

题意翻译
有一颗n个结点且有n-1条边的无根树

一只熊可以从当前节点可以跳到任何与当前节点距离不超过k的节点

定义f(u,v)f(u,v)为熊从u点到v点所需的最少跳跃次数，那么，对于树上的所有点对(u,v)(u,v)，f(u,v)f(u,v)的总和是多少。

注意：

(u,v)(u,v)和(v,u)(v,u)视作同一个点对，只计算一次答案

输入输出样例
输入 #1复制
6 2
1 2
1 3
2 4
2 5
4 6
输出 #1复制
20
输入 #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
输出 #2复制
114
输入 #3复制
3 5
2 1
3 1
输出 #3复制
3
说明/提示
In the first sample, the given tree has 66 vertices and it’s displayed on the drawing below. Limak can jump to any vertex within distance at most 22 . For example, from the vertex 55 he can jump to any of vertices: 11 , 22 and 44 (well, he can also jump to the vertex 55 itself).

There are pairs of vertices (s,t)(s,t) such that s<t . For 55 of those pairs Limak would need two jumps: (1,6),(3,4),(3,5),(3,6),(5,6)(1,6),(3,4),(3,5),(3,6),(5,6) . For other 1010 pairs one jump is enough. So, the answer is 5·2+10·1=205⋅2+10⋅1=20 .

In the third sample, Limak can jump between every two vertices directly. There are 33 pairs of vertices (s<t) , so the answer is 3·1=33⋅1=3 .

思路

考虑树形dp
设f[u][i]为u的子树中与u的距离为i的点的答案
对于i>0，可以由儿子转移到父亲
对于i=0，也就说点u的答案，就要分类讨论

1. 对于距离<k的点，贡献为1
2. 对于距离>k的点，可以先跳k格，然后由那一个点转移

转移方程长这样：

代码

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=2e5+77;
int n,k,sz[N];
vector<int> g[N];
ll f[N][5];
void dfs1(int u,int p)
{
	sz[u]=1;
	for(int v:g[u])
		if(v!=p)
		{
			dfs1(v,u);
			sz[u]+=sz[v],f[u][0]+=f[v][k-1]+sz[v];
			for(int i=1; i<k; i++) f[u][i]+=f[v][i-1];
		}
}
void dfs2(int u,int p)
{
	ll t[5]={0};
	if(p)
	{
		t[0]=f[p][0]-f[u][k-1]-sz[u];
		for(int i=1; i<k; i++) t[i]=f[p][i]-f[u][i-1];
	}
	f[u][0]+=t[k-1]+n-sz[u];
	for(int i=1; i<k; i++) f[u][i]+=t[i-1];
	for(int v:g[u]) if(v!=p) dfs2(v,u);
}
int main()
{
	scanf("%d%d",&n,&k);
	for(int i=1; i<n; i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		g[u].push_back(v); g[v].push_back(u);
	}
	dfs1(1,0); dfs2(1,0);
	ll ans=0;
	for(int i=1; i<=n; i++) ans+=f[i][0];
	printf("%lld",ans/2);
}