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;
}