luogu P3047 [USACO12FEB]Nearby Cows G 换根dp

#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;
const int N = 2e5 + 10;
int h[N], e[N], ne[N], idx;
int ans[N];
int a[N];
int f[N][25];
int n, m;
void add(int a, int b)
{
	e[idx] = b;
	ne[idx] = h[a];
	h[a] = idx++;
} 
void cut(int now, int pre)
{
	for(int i = 1; i <= m; i ++)
		f[now][i] -= f[pre][i - 1];
}
void link(int now, int pre)
{
	for(int i = 1; i <= m; i ++)
		f[now][i] += f[pre][i - 1];
}
void dfs1(int now,int fa)
{
	f[now][0] = a[now];
	for(int i = h[now]; ~i; i = ne[i])
	{
		int v = e[i];
		if(v == fa)
			continue;
		dfs1(v, now);
		for(int j = 1; j <= m; j ++)
			f[now][j] += f[v][j-1];
	}
}
void dfs2(int now, int fa)
{
	//与当前点的距离在m以内的 
	for(int i = 0; i <= m; i ++)
		ans[now] += f[now][i];
	for(int i = h[now]; ~i; i = ne[i])
	{
		int v = e[i];
		if(v == fa)
			continue;
		//先把v所在的子树中，到v的距离小于等于m-1的砍掉
		//那么也就是在now中，砍掉v所在子树中，到now的距离小于m的 
		cut(now, v);
		//然后v再加上，除了v所在子树中，到now距离小于m-1
		//也就是到v距离为m的 
		link(v, now);
		
		dfs2(v, now);
		//恢复现场 
		cut(v, now);
		link(now, v);
	}
}
int main()
{
	memset(h, -1, sizeof h);
	cin >> n >> m;
	for(int i = 1; i < n; i ++)
	{
		int a, b;
		cin >> a >> b;
		add(a, b);
		add(b, a);
	}
	for(int i = 1; i <= n; i ++)
		cin >> a[i];
	dfs1(1, -1);
	dfs2(1, -1);
	for(int i = 1; i <= n; i ++)
		cout << ans[i] << endl;
	return 0;
}