[USACO12FEB]Nearby Cows G（换根dp）

一句话题意：给你一棵 \(n\) 个点的树，点带权，对于每个节点求出距离它不超过 \(k\) 的所有节点权值和 \(m_i\) 。

\(1 \le n \le 10^5\)

定睛一看这就是今年省选B卷D1T2的60pts数据嘛。

k的范围很小，可以用\(O(nk)\)的算法水过去。其实就是换根dp。

令\(dp[u][k]\)代表u到其子树内距离为k的点的权值和，有\(dp[u][k]+=dp[v][k-1]\)。

然后考虑怎么从父亲扩展到儿子，其实就是一步容斥：\(f[v][k]=dp[v][k]+f[u][k-1]-dp[v][k-2]\)。

#include 
#include 
#include 
#include 
#include 
using namespace std;
typedef long long ll;
const int N = 100010;
const int K = 22;
const int inf = 0x3f3f3f3f;
template  void read(T &x) {
	T w = 1;
	char ch = getchar();
	for (; !isdigit(ch); ch = getchar()) if (ch == '-') w = -1;
	for (x = 0; isdigit(ch); ch = getchar()) x = x * 10 + ch - '0';
	x *= w;
}
struct node{
	int pre, to;
}edge[N << 1];
int head[N], tot;
int n, k;
int dp[N][K], f[N][K];
void add(int u, int v) {
	edge[++tot] = node{head[u], v};
	head[u] = tot;
}
void dfs1(int x, int fa) {
	for (int i = head[x]; i; i = edge[i].pre) {
		int y = edge[i].to;
		if (y == fa) continue;
		dfs1(y, x);
		for (int j = 1; j <= k; j++) {
			dp[x][j] += dp[y][j - 1];
		}
	}
}
void dfs2(int x, int fa) {
	f[x][0] = dp[x][0];
	for (int i = head[x]; i; i = edge[i].pre) {
		int y = edge[i].to;
		if (y == fa) continue;
		for (int j = 1; j <= k; j++) {
			if (j > 1) f[y][j] = dp[y][j] + (f[x][j - 1] - dp[y][j - 2]);
			else f[y][j] = dp[y][j] + (f[x][j - 1]);
		}
		dfs2(y, x);
	}
}
int main() {
	read(n); read(k);
	for (int i = 1, u, v; i < n; i++) {
		read(u); read(v);
		add(u, v);
		add(v, u);
	}
	for (int i = 1; i <= n; i++) read(dp[i][0]);
	dfs1(1, 0);
	for (int i = 1; i <= k; i++) f[1][i] = dp[1][i];
	dfs2(1, 0);
	for (int i = 1; i <= n; i++) {
		for (int j = 0; j <= k; j++) {
			f[i][j] += f[i][j - 1];
		}
		printf("%d\n", f[i][k]);
	}
	return 0;
}