【洛谷-P3047】[USACO12FEB]Nearby Cows G • TonyYin's Blog

题意#

给定一棵 $n$ 个点的数，点带权。对于每个节点 $i$ ，求出距离 $i$ 不超过 $k$ 的所有节点的权值和。

$1\leq n\leq 10^5$ ， $1\leq k\leq 20$ ， $0\leq c_i\leq 1000$ .

分析#

考虑树形DP。

先任意钦定 $\rm{root}$ .

设 $f_{i, j}$ 表示在 $i$ 的子树中，距离 $i$ 小于等于 $j$ 的节点的权值和。

可以用一次 $\rm {DFS}$ 求得 $f$ ，转移如下：

f_{x, i}=c_x+\sum_{v} f_{v, i-1}

设 $g_{i, j}$ 表示所有距离 $i$ 距离小于等于 $j$ 的节点的权值和，所以 $\operatorname{Ans}_i = g_{i,k}$ .

再用一次 $\rm{DFS}$ 求得 $g$ ，转移如下：

g_{x,i} = \left\{\begin{array}{lr} f_{x, i}+\sum\limits_{v}(g[\operatorname{fa}][i-1]-f_{x, i-2}), &x=1\\ f_{x, i}, &x\neq 1 \end{array}\right.

注意当 $i=1$ 时，应忽略 $f_{x, i-2}$ 项，否则会数组越界。

代码#

const int MAXN = 1e5 + 10, MAXK = 25;
int n, k, c[MAXN];
vector<int> g[MAXN];
int f[MAXN][MAXK];
void dfs1(int x, int fa) {
	for(int i = 0; i <= k; i++) f[x][i] = c[x];
	for(int v : g[x]) {
		if(v == fa) continue;
		dfs1(v, x);
		for(int i = 1; i <= k; i++) f[x][i] += f[v][i - 1];
	}
}
void dfs2(int x, int fa) {
	if(x != 1) {
		for(int i = k; i >= 2; i--)
			f[x][i] += f[fa][i - 1] - f[x][i - 2];
		f[x][1] += f[fa][0];
	}
	for(int v : g[x]) {
		if(v == fa) continue;
		dfs2(v, x);
	}
}
int main() {
	n = read(), k = read();
	for(int i = 1; i < n; i++) {
		int u = read(), v = read();
		g[u].push_back(v), g[v].push_back(u);
	}
	for(int i = 1; i <= n; i++) c[i] = read();
	dfs1(1, 0); dfs2(1, 0);
	for(int i = 1; i <= n; i++) cout << f[i][k] << "\n";
	return 0;
}

cpp