【bzoj2591】[Usaco 2012 Feb]Nearby Cows 树形dp

题目描述

Farmer John has noticed that his cows often move between nearby fields. Taking this into account, he wants to plant enough grass in each of his fields not only for the cows situated initially in that field, but also for cows visiting from nearby fields. Specifically, FJ's farm consists of N fields (1 <= N <= 100,000), where some pairs of fields are connected with bi-directional trails (N-1 of them in total). FJ has designed the farm so that between any two fields i and j, there is a unique path made up of trails connecting between i and j. Field i is home to C(i) cows, although cows sometimes move to a different field by crossing up to K trails (1 <= K <= 20). FJ wants to plant enough grass in each field i to feed the maximum number of cows, M(i), that could possibly end up in that field -- that is, the number of cows that can potentially reach field i by following at most K trails. Given the structure of FJ's farm and the value of C(i) for each field i, please help FJ compute M(i) for every field i.

FJ发现他的牛经常跑到附近的草地去吃草，FJ准备给每个草地种足够的草供这个草地以及附近草地的奶牛来吃。FJ有N个草地(1<=N<=100000)，有N-1条双向道路连接这些草地，FJ精心设计了这些道路使每两个草地有且仅有一条简单路径连接。第i个草场有Ci头牛，有时候奶牛会走过K条道路到其他草地吃草。FJ想知道每个草场最多可能有的奶牛数量Mi，即所有走过K条道路后可能到达i的奶牛总数。

输入

* Line 1: Two space-separated integers, N and K.

* Lines 2..N: Each line contains two space-separated integers, i and j (1 <= i,j <= N) indicating that fields i and j are directly connected by a trail.

* Lines N+1..2N:

Line N+i contains the integer C(i). (0 <= C(i) <= 1000)

输出

* Lines 1..N: Line i should contain the value of M(i).

样例输入

6 2
5 1
3 6
2 4
2 1
3 2
1
2
3
4
5
6

样例输出

15
21
16
10
8
11

---

题解

树形dp

由于直接推出答案比较困难，所以我们可以现在子树中寻找答案，然后再求父树中。

设f[x][k]表示x子树中到x的距离为k的的点的点权之和，

显然有f[x][k]=∑f[to[i]][k-1]。

再设g[x][k]表示整棵树中到x的距离为k的点的点权之和。

首先有g[1][k]=f[1][k]，g[x][0]=v[x]，剩下的就只能从上到下dp了。

我们已经知道子树中的答案，差的就是父树。

而如果父树中的点到to[i]的距离为k，那么这些点到x的距离必然为k-1.

于是我们就可以通过g[x][k-1]推出g[to[i]][k]。

然而这样当k≥2时会重复计算子树中的某些点，这些点到x的为k-1，那么到to[i]的距离必然为k-2。

这样就再减去f[to[i]][k-2]即可。

最终状态转移方程就为g[to[i]][k]=f[to[i]][k]+g[x][k-1]-f[to[i]][k-2]。

最后答案就是∑g[i][j] (0≤j≤k)

#include <cstdio>
#include <algorithm>
using namespace std;
int f[100010][25] , g[100010][25] , c[100010] , head[100010] , to[200010] , next[200010] , cnt , k;
void add(int x , int y)
{
	to[++cnt] = y;
	next[cnt] = head[x];
	head[x] = cnt;
}
void dp1(int x , int fa)
{
	int i , j;
	f[x][0] = c[x];
	for(i = head[x] ; i ; i = next[i])
	{
		if(to[i] != fa)
		{
			dp1(to[i] , x);
			for(j = 1 ; j <= k ; j ++ )
				f[x][j] += f[to[i]][j - 1];
		}
	}
}
void dp2(int x , int fa)
{
	int i , j;
	for(i = head[x] ; i ; i = next[i])
	{
		if(to[i] != fa)
		{
			g[to[i]][0] = c[to[i]];
			for(j = 1 ; j <= k ; j ++ )
			{
				g[to[i]][j] = f[to[i]][j] + g[x][j - 1];
				if(j >= 2) g[to[i]][j] -= f[to[i]][j - 2];
			}
			dp2(to[i] , x);
		}
	}
}
int main()
{
	int n , i , j , x , y , ans;
	scanf("%d%d" , &n , &k);
	for(i = 1 ; i < n ; i ++ )
		scanf("%d%d" , &x , &y) , add(x , y) , add(y , x);
	for(i = 1 ; i <= n ; i ++ )
		scanf("%d" , &c[i]);
	dp1(1 , 0);
	for(i = 0 ; i <= k ; i ++ )
		g[1][i] = f[1][i];
	dp2(1 , 0);
	for(i = 1 ; i <= n ; i ++ )
	{
		ans = 0;
		for(j = 0 ; j <= k ; j ++ )
			ans += g[i][j];
		printf("%d\n" , ans);
	}
	return 0;
}