[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\ast dep[lca]$，$f(a,b)=(dis(a,b)-k)modk$,转移的时候我们就可以枚举$u,v$中dp状态第二维$i,j$，这些点之间的距离$dis=(i+j-dep[u])modk$（因为$u,v$的$lca$是$u$）单个点对需要的代价$f=（dis-k)modk$，由乘法原理得总的代价为$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;
}