Codeforces VK Cup 2017 - Round 1 C 树形dp 双解法

题意: 每次可以跳跃不超过k的距离 求树上所有点对之间的最短距离和。

解法一:

题目所求即为 所有路径的长度/k向上取整的和  --> sum( (dis(i, j)-1) / k + 1 )

考虑对所有对k取余不为0的路径长度进行向上取整

即加上每条路径长度的(k-res)

那么我们需要求出所有树上路径长度的和以及所有路径长度对k的余数

令dp[i][j]为 i的子树中到根节点的距离%k为j的节点数 dp的过程中累计答案即可

#include <iostream> /// cf VK Cup 2017 - Round 1 C 每次可以跳跃不超过k的距离 求树上所有点对之间的最短距离和 O(n * k^2)
#include <cstring>
using namespace std; // 即为求 所有路径的长度/k向上取整的和  --> sum( (dis(i, j)-1) / k + 1 )

const int N=200005,mod=1e9+7;
struct pp
{
	int to;
	int old;
}edge[N<<1];
int sz[N],dep[N],dp[N][6]; // dp[i][j] --> i的子树中到根节点的距离%k为j的节点数
int newly[N],cnt,n,k;
long long ans;
void add(int u, int v)
{
	edge[cnt]={v, newly[u]};
	newly[u]=cnt++;
}
void dfs(int x, int fa)
{
	sz[x]=1;
	dep[x]=dep[fa]+1;
	dp[x][dep[x]%k]=1;
	for(int i=newly[x]; ~i; i=edge[i].old)
	{
		int son=edge[i].to;
		if(son==fa)
			continue;
		dfs(son, x);
		sz[x]+=sz[son];
		for(int p=0; p<k; p++)
			for(int q=0; q<k; q++)
			{
				int dis=(p+q-2*dep[x]%k+k)%k; // 前面子树中的点 到 当前子树中的点 的最短距离对k取模
				int res=(k-dis)%k; // 要加上比k缺少的
				ans+=1ll*res*dp[x][p]*dp[son][q]; // 相当于对原式向上取整了
			}
		for(int j=0; j<k; j++)
			dp[x][j]+=dp[son][j];
	}
	ans+=1ll*sz[x]*(n-sz[x]);
}

int main()
{
	memset(newly,-1,sizeof(newly));
	scanf("%d%d", &n, &k);
	for(int i=1; i<n; i++)
	{
		int x,y; scanf("%d%d", &x, &y);
		add(x, y),add(y, x);
	}
	dfs(1, 0);
	cout << ans/k << '\n';
	return 0;
}

时间复杂度 O(n * k^2) 适用于k较小的情况

解法二:

考虑换根dp

令down[i][j]为i的子树中到节点i距离%k为j的节点跳到i节点的答案和

注意第二遍dfs去除儿子节点方向的贡献时若 dp[fa][j] 的 j==0 则要额外减去sz[son]

#include <iostream> /// cf VK Cup 2017 - Round 1 C 每次可以跳跃不超过k的距离 求树上所有点对之间的最短距离和 O(n * k)
#include <cstring>
using namespace std; // 换根dp

const int N=200005;
struct pp
{
	int to;
	int old;
}edge[N<<1];
long long down[N][6],sz[N]; // down[i][j] --> i的子树中到节点i距离%k为j的节点的子树答案和
long long dp[N][6];
int newly[N],cnt,n,k;
long long ans;
void add(int u, int v)
{
	edge[cnt]={v, newly[u]};
	newly[u]=cnt++;
}

void dfs_1(int x, int fa)
{
	sz[x]=1;
	for(int i=newly[x]; ~i; i=edge[i].old)
	{
		int son=edge[i].to;
		if(son==fa)
			continue;
		dfs_1(son, x);
		down[x][0]+=down[son][k-1]+sz[son];
		for(int j=1; j<k; j++)
			down[x][j]+=down[son][j-1];
		sz[x]+=sz[son];
	}
}

void dfs_2(int x, int fa)
{
	for(int i=newly[x]; ~i; i=edge[i].old)
	{
		int son=edge[i].to;
		if(son==fa)
			continue;
		dp[son][0]=down[son][0]+dp[x][k-1]-down[son][(k-2+k)%k]-(k==1 ? sz[son] : 0)+(n-sz[son]); // k==1时同下
		for(int j=1; j<k; j++) // 当j==1时 dp[x][0]来自down[son]的贡献会额外加上son子树中的节点转移的步数
			dp[son][j]=down[son][j]+(dp[x][j-1]-down[son][(j-2+k)%k]-(j==1 ? sz[son] : 0));
		dfs_2(son, x);
	}
}

int main()
{
	memset(newly,-1,sizeof(newly));
	scanf("%d%d", &n, &k);
	for(int i=1; i<n; i++)
	{
		int x,y; scanf("%d%d", &x, &y);
		add(x, y),add(y, x);
	}
	dfs_1(1, 0);
	for(int i=0; i<k; i++)
		dp[1][i]=down[1][i];
	dfs_2(1, 0);
	for(int i=1; i<=n; i++)
		ans+=dp[i][0];
	cout << ans/2 << '\n';
	return 0;
}

时间复杂度 O(n * k)