P3047 [USACO12FEB]Nearby Cows G（树形DP）

题意：

给出一棵树。

对于每个点，求出距离它不超过k的所有点的点权和。

题解：

定义\(f(i,j)\)表示第i个点的j范围内的点权和。

定义\(gg(i,j)\)表示第i个点向下j范围内的点权和。

\(gg(i,j)=c_i+\sum_{v \in son_u}gg(v,j-1)\)

\(f(i,j)=f(fa_i,j-1)-gg(i,j-2)+gg(i,j)\)

第一遍DFS求出所有的\(gg(i,j)\)。

第二遍DFS求出所有的\(f(i,j)\)。

#include<bits/stdc++.h>
using namespace std;
const int maxn=2e5+100;
vector<int> g[maxn];
int a[maxn];
int n,k;
int gg[maxn][30];
int f[maxn][30];
void dfs1 (int x,int pre) {
	gg[x][0]=a[x];
	for (int y:g[x]) {
		if (y==pre) continue;
		dfs1(y,x);
		for (int j=1;j<=k;j++) gg[x][j]+=gg[y][j-1];
	}
}
void dfs2 (int x,int pre) {
	f[x][0]=gg[x][0];
	for (int j=1;j<=k;j++) f[x][j]=f[pre][j-1]+gg[x][j];
	for (int j=2;j<=k;j++) if (pre) f[x][j]-=gg[x][j-2];
	for (int y:g[x]) {
		if (y==pre) continue;
		dfs2(y,x);
	}
}
int main () {
	cin>>n>>k;
	for (int i=1;i<n;i++) {
		int x,y;
		cin>>x>>y;
		g[x].push_back(y);
		g[y].push_back(x);
	}
	for (int i=1;i<=n;i++) cin>>a[i];
	dfs1(1,0);
	for (int i=1;i<=n;i++) for (int j=1;j<=k;j++) gg[i][j]+=gg[i][j-1];
	dfs2(1,0);
	for (int i=1;i<=n;i++) printf("%d\n",f[i][k]);
}