791D Bear and Tree Jumps

D. 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 ofn vertices, numbered 1 through n.

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

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

Input

The first line of the input contains two integersn and k (2 ≤ n ≤ 200 000,1 ≤ k ≤ 5) — the number of vertices in the tree and the maximum allowed jump distance respectively.

The next n - 1 lines describe edges in the tree. Thei-th of those lines contains two integers a_i and b_i (1 ≤ a_i, b_i ≤ n) — the indices on vertices connected withi-th edge.

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

Output

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

如果小熊每次能跳跃的距离为1，那么问题变为求树上任意两点之间距离之和。
对于每一条边sum1和sum2分别表示边的左右点数，ans=各边的sum1*sum2之和即为答案。
而本题最大跳跃距离为k，答案变为(ans+sum)/k。sum为每一条边需要多走x步才能整除k的x之和。
树上任意两点间距离len=depth[x1]+depth[y1]-2*depth[f]，f表示点x1和点y1的最近公共祖先。
计算sum的方法：dp[i][j]表示到i点的距离对k取摸为j的点的总数。
则对于任意两点a和b，dis需要满足(len[a][b]+dis)%k==0。
每当搜索到一个点时，用O(k^2)的方法枚举所有情况。

#include<stdio.h>
#include<algorithm>
#include<string.h>
#include<vector>
#define ll long long
using namespace std;
const int maxm = 200005;
vector<int>v[maxm];
ll dp[maxm][5] = { 0 }, a[maxm] = { 0 }, ans = 0;
int k, n;
void dfs(int x, int pre, int num);
int main()
{
	int  i, j, sum, x, y;
	scanf("%d%d", &n, &k);
	for (i = 1;i < n;i++)
	{
		scanf("%d%d", &x, &y);
		v[x].push_back(y);
		v[y].push_back(x);
	}
	dfs(1, -1, 0);
	printf("%lld\n", ans / k);
	return 0;
}
void dfs(int x, int pre, int num)
{
	dp[x][num%k] = 1;
	a[x] = 1;
	for (int i = 0;i < v[x].size();i++)
	{
		int xx = v[x][i];
		if (xx == pre)
			continue;
		dfs(xx, x, num + 1);
		for (int j = 0;j < k;j++)
		{
			for (int r = 0;r < k;r++)
			{
				int dis = (j + r - num * 2) % k;
				int rev = (k - dis) % k;
				ans += rev*dp[x][j] * dp[xx][r];
			}
		}
		a[x] += a[xx];
		for (int j = 0;j < k;j++)
			dp[x][j] += dp[xx][j];
		ans += (n - a[xx])*a[xx];
	}
}