CF771C Bear and Tree Jumps

题目大意：

给定一棵有 $n$ 个节点的树，要你统计 $\sum _{1\le x\le y\le n}dist(x,y)/k$ ($dist(x,y)$ 表示 $x$ 到 $y$ 的距离)

$n\le 2\times {10}^5,k\le 5$

解法：

一道换根 $dp$ 套路题。

首先看到树上统计问题，考虑树形 $dp$，那么我们设 $g(u)$ 为以 $u$ 为根节点时子树的答案。

此时我们发现，这个状态定义十分的不可做，无法计算贡献和转移。

观察到 $k$ 的值域很小，所以多添加一维距离 $i$。具体的状态定义就变为了 $g(u,i)$ 表示在以 $u$ 为根的子树，距离为 $i$ 的节点的产生的贡献和。

这个时候就很好转移了，直接有：

• $g(u,0)=\sum g(v,k-1)+siz[v]$

• $g(u,i)=\sum g(v,i-1)|i>0$

但是这统计的只是子树中的贡献，没有统计子树外的，那么自然而然的就想到了换根 $dp$。

设 $f(u,i)$ 为整棵树中，与 $u$ 距离为 $i$ 的节点产生的贡献和。

显然的可以推出公式了。

code:

#include<bits/stdc++.h>
#define int long long 
using namespace std;
 
const int N=2e5+10;
int n,k,g[N][6],f[N][6],siz[N],ans,val[6];
struct edge{
	int v,next;
}edges[N*2];
int head[N],idx;
void add_edge(int u,int v){
	idx++;
	edges[idx].v=v;
	edges[idx].next=head[u];
	head[u]=idx;
	return;
}
void dfs1(int u,int fa){
	siz[u]=1;
	for(int i=head[u];i;i=edges[i].next){
		int v=edges[i].v;
		if(v==fa)continue;
		dfs1(v,u);
		g[u][0]+=siz[v]+g[v][k-1];
		for(int j=1;j<k;j++)g[u][j]+=g[v][j-1];
		siz[u]+=siz[v];
	}
	return;
}
void dfs2(int u,int fa){
	if(u==1){for(int i=0;i<=k;i++)f[u][i]=g[u][i];}
	else{
		val[0]=f[fa][0]-g[u][k-1]-siz[u];
		for(int i=1;i<k;i++)val[i]=f[fa][i]-g[u][i-1];
		f[u][0]=g[u][0]+val[k-1]+n-siz[u];
		for(int i=1;i<k;i++)f[u][i]=g[u][i]+val[i-1];
	}
	for(int i=head[u];i;i=edges[i].next){
		int v=edges[i].v;
		if(v!=fa)dfs2(v,u);
	}
	return;
}
 
signed main(){
	std::ios::sync_with_stdio(false);
	cin.tie(0);cout.tie(0);
	cin>>n>>k;
	for(int i=1;i<n;i++){
		int x,y;
		cin>>x>>y;
		add_edge(x,y);add_edge(y,x);
	}
	dfs1(1,0);dfs2(1,0);
	for(int i=1;i<=n;i++)ans+=f[i][0];
	cout<<ans/2;
	
	return 0;
}