[USACO12FEB]Nearby Cows G

(large{题目链接})
(\)
(Large extbf{Solution: } large{普普通通的一道换根dp。\设f_i表示从i点出发的答案，d_{i,j}表示i向下距离为j的点权和，up_{i,j}表示i向上距离为j的点权和，然后大力转移即可。\然后我发现其实不用up数组，直接f和d容斥一下就行。})
(\)
(Large extbf{Code: })

#include <bits/stdc++.h>
using namespace std;

const int N = 1e5 + 5;

int n, k, a[N], f[N], d[N][25], up[N][25];
vector <int> v[N];

int read(int &x) {
	int flg = 1; x = 0;
	char c = getchar();
	while (!isdigit(c)) { if (c == '-') flg = -1; c = getchar(); }
	while (isdigit(c)) x = x * 10 + c - '0', c = getchar();
	x *= flg;
}

void dfs(int x, int fa) {
	for (int i = 0; i < v[x].size(); ++i) {
		int u = v[x][i];
		if (u == fa) continue;
		dfs(u, x);
		for (int j = 1; j <= k; ++j) d[x][j] += d[u][j - 1];
	}
}

void dfs2(int x, int fa) {
	up[x][0] = a[x];
	for (int i = 0; i < v[x].size(); ++i) {
		int u = v[x][i];
		if (u == fa) continue;
		f[u] = f[x] - up[x][k] + d[u][k] - d[x][k] + d[u][k - 1];
		for (int j = 1; j <= k; ++j) up[u][j] = up[x][j - 1] + d[x][j - 1] - d[u][j - 2];
		up[u][1] -= d[x][0];
		dfs2(u, x);
	}
}

int main() {
	read(n), read(k);
	int x, y;
	for (int i = 2; i <= n; ++i) read(x), read(y), v[x].push_back(y), v[y].push_back(x);
	for (int i = 1; i <= n; ++i) read(a[i]);
	for (int i = 1; i <= n; ++i) d[i][0] = a[i];
	dfs(1, 0);
	for (int i = 0; i <= k; ++i) f[1] += d[1][i];
	dfs2(1, 0);
	for (int i = 1; i <= n; ++i) printf("%d
", f[i]);
	return 0; 
}