贡献法之求树上上任意两点间的路径之和

一图胜千言

开始造火箭

虽然我们可以求出，总共的dis(i,j)，但分散到每一个小dis(i,j)，由于存在向上取整操作，我们需要求出将每一个小dis(i,j)给补成k的倍数的补数之和！

此处我们采用树形dp。

dp[u][j]表示以u的子节点到根节点root的距离对k取余的值为j的点的个数

我们如何求树上任意两点的距离呢？

dis(a,b) = depth(a) + depth(b) - 2 * depth(lca(a,b))

那么，dis(a,b)的补数 **(k - dis(a,b)) % k **

上代码

#include<iostream>
#include<cstring>
#include<algorithm>
#define ll long long
const int N = 2e5 + 5; 
using namespace std;
int n,k;
int h[N],e[N * 2],ne[N * 2],idx;
ll dp[N][5],s[N];
ll ans;
void add(int a,int b)
{
	e[idx] = b;
	ne[idx] = h[a];
	h[a] = idx++;
}
void dfs(int u,int fa,int depth)
{
	dp[u][depth % k] = 1;s[u] = 1;
	for(int i = h[u];~i;i = ne[i])
	{
		int v = e[i];
		if(v == fa) continue;
		dfs(v,u,depth + 1);
		for(int j = 0;j < k;++j)
		{
			for(int r = 0;r < k;++r)
			{
				int dis = (j + r - 2 * depth) % k;
				int bushu = (k - dis) % k;
				ans += bushu * dp[u][j] * dp[v][r];
			}
		}
		s[u] += s[v];
		for(int j = 0;j < k;++j) dp[u][j] += dp[v][j];
		ans += (n - s[v]) * s[v];
	}
}
int main()
{
	cin >> n >> k;
	memset(h,-1,sizeof(h));
	for(int i = 1;i < n;++i)
	{
		int a,b;cin >> a >> b;
		add(a,b);add(b,a);
	}
	dfs(1,-1,0);
	cout << ans / k << '\n';
	return 0;
}