[CF791D]Bear and Tree Jumps 树上DP

题面：

给定一颗大小为\(n\)的树，给定一个k，求\(\sum_{i=1}^n \sum_{j=1}^n \lceil \frac{dis(i,j)}{k} \rceil\)

范围&性质：\(1\le n \le 2\times 10^5,1\le k\le 5\)

分析：

朴素暴力做法：

对于每一个点dfs一遍，统计答案，复杂度\(O(n^2)\)

正解：（前排膜拜 Orz）

先考虑k=1怎么做，相当于求树上任意两点间路径和，直接枚举每条边，求出左右两部分内点的个数\(num1,num2\)，相乘就好了，复杂度\(O(n)\)

那怎么向k>1扩展呢，问题在于k>1时路径长度除以k会有余数，那我们就考虑把余数补齐，答案就变成了 \(\frac{ans+\sum f(i,j)}{k}\)，其中\(f(i,j)\)表示将\(dis(i,j)\)补成k的倍数需要的代价，这样就可以消掉上取整符号的影响了

那么问题就转化成了如何求\(\sum f(i,j)\),这时候就利用到k很小这一性质，我们记\(dp[i][j]\)表示u的子树内到根节点距离模\(k\)余\(j\)的点的个数，对于树上任意两点\(a,b\)而言，他们之间的距离\(dis(a,b)=dep[a]+dep[b]-2*dep[lca]\)，\(f(a,b)=(dis(a,b)-k)mod \ k\),转移的时候我们就可以枚举\(u,v\)中dp状态第二维\(i,j\)，这些点之间的距离\(dis=(i+j-dep[u]) mod\ k\)（因为\(u,v\)的\(lca\)是\(u\)）单个点对需要的代价\(f=（dis-k)mod\ k\)，由乘法原理得总的代价为\(f(i,j)\times dp[u][i]\times dp[v][j]\)

代码：

#include<bits/stdc++.h>

using namespace std;

namespace zzc
{
	const int maxn = 200005;
	int cnt=0,head[maxn];
    long long dp[maxn][5],siz[maxn],ans=0;
    int k,n;
	
	struct edge
    {
    	int to,nxt;
	}e[maxn<<1];
	
	void add(int u,int v)
	{
		e[++cnt].to=v;
		e[cnt].nxt=head[u];
		head[u]=cnt;
	}
	
	void dfs(int u,int fa,int dep)
	{
		dp[u][dep%k] = 1;siz[u]=1;
		for (int i=head[u];i;i=e[i].nxt)
		{
			int v=e[i].to;
			if (v==fa) continue;
			dfs(v,u,dep+1);
			for (int j = 0;j < k;j++)
			{
				for (int r = 0;r < k;r++)
				{
					int dis = (j + r - dep * 2) % k;
					int rev = (k - dis) % k;
					ans += rev*dp[u][j] * dp[v][r];
				}
			}
			siz[u]+=siz[v];
			for (int j = 0;j < k;j++) dp[u][j] += dp[v][j];
			ans += (n - siz[v])*siz[v];
		}
	}

	
	void work()
	{
		scanf("%d%d",&n,&k);
		for (int i=1,x,y;i < n;i++)
		{
			scanf("%d%d",&x,&y);
			add(x,y);add(y,x);
		}
		dfs(1, -1, 0);
		printf("%lld\n",ans/k);
	}
	
}

int main()
{
	zzc::work();
	return 0;
}