C. Bear and Tree Jumps(树形dp+换根)

题目

题意：

    给定一棵树，再给一个k，要求计算树上任意两点的距离除以k(向上取整)的和。
     $2 ≤ n ≤ 200 000, 1 ≤ k ≤ 5$

分析：

    因为是树上问题，所以自然就想到子节点与父节点的转移。因为k比较小，所以我们维护一个num[i][j]，表示到i节点的子树中到i的距离模k为j的节点数。那么对于子节点到父节点的转移，其实只要加上模k为0的那些点，到父节点需要加1的代价和当前子节点到父节点的代价，其余的只要用来更新父节点的num数组即可。这样我们就可以算出根节点的贡献。
    由于要算其他的点的贡献，所以我们需要换根。换根的话就是将父节点的num数组减去当前子节点对它的贡献，然后子节点加上父节点的num数组，然后再算上父节点num数组模k为0的点的贡献和父节点的贡献(原先的父节点现在变成了子节点)。记得回溯时还原！！！

#include <iostream>
#include <vector> 
using namespace std;

typedef long long ll;

ll dp[200005],num[200005][5];
int n,k;
vector<int> g[200005]; 

void dfs1(int x,int fa)
{
	if( g[x].size() == 1 && g[x][0] == fa )
	{
		return;
	}
	for (int i = 0; i < g[x].size(); i++)
	{
		int t = g[x][i];
		if( t == fa ) continue;
		dfs1(t,x);
		for (int j = 0; j < k; j++)
		{
			num[x][j] += num[t][(j-1+k)%k];
		}
		num[x][1%k] ++;
		dp[x] += dp[t] + 1 + num[t][0];
	}
}

void dfs2(int x,int fa)
{
	for (int i = 0; i < g[x].size(); i++)
	{
		int t = g[x][i];
		if( t == fa ) continue;
		for (int j = 0; j < k; j++)
		{
			num[x][j] -= num[t][(j-1+k)%k];
		}
		num[x][1%k] --;
		dp[t] += dp[x] - dp[t] - 1 - num[t][0];
		dp[t] += 1 + num[x][0];
		for (int j = 0; j < k; j++)
		{
			num[t][j] += num[x][(j-1+k)%k];
		}
		num[t][1%k] ++;
		dfs2(t,x);
		num[t][1%k] --;
		for (int j = 0; j < k; j++)
		{
			num[t][j] -= num[x][(j-1+k)%k];
		}
		num[x][1%k] ++;
		for (int j = 0; j < k; j++)
		{
			num[x][j] += num[t][(j-1+k)%k];
		}
	}
}

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	cin >> n >> k;
	for (int i = 1; i < n; i++)
	{
		int x,y;
		cin >> x >> y;
		g[x].push_back(y);
		g[y].push_back(x);
	}
	dfs1(1,0);
	dfs2(1,0);
	ll ans = 0;
	for (int i = 1; i <= n; i++)
	{
		ans += dp[i];
	}
	cout << ans / 2 << '\n';
	return 0;
}