[USACO12FEB] Nearby Cows G

题目描述

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

输入格式

第一行两个正整数 $n,k$。
接下来 $n-1$ 行，每行两个正整数 $u,v$，表示 $u,v$ 之间有一条边。
最后 $n$ 行，每行一个非负整数 $c_i$，表示点权。

输出格式

输出 $n$ 行，第 $i$ 行一个整数表示 $m_i$。

样例 #1

样例输入 #1

样例输出 #1

【数据范围】
对于 $100\%$ 的数据：$1 \le n \le 10^5$，$1 \le k \le 20$，$0 \le c_i \le 1000$

解决方案

树上DP+两次DFS

首先对于某一个结点$u$，考虑答案的组成部分：
1、$u$的子树中距离不超过$k$的所有点的权值和；
2、不在$u$的子树中的距离不超过$k$的所有点的权值和；

第一部分用一个简单的树形$DP$（第一次$DFS$）即可：设$f[i][j]$表示以$i$为根的子树中距离为$j$的所有节点的权值和。则$f[i][j]$就等于距离$i$的所有子结点距离不超过$j-1$的权值和，即状态转移方程为$$f[i][j] = \sum f[k][j-1]$$，其中$k \in i$的子节点。
经过第一次$DFS$后，我们就求出了对于每一个节点$i$，其子树中所有满足要求的答案。然而题目的要求不止子树，因此还需要对非$i$的子树中的节点。
设$i$的父节点是$fa$，$d[i][j]$表示距离节点$i$距离为$j$的所有节点的权值和。首先初始化$d[i][j] == f[i][j]$，然后$d[i][j] += d[fa][j - 1]$。但是这样会有重复，因为距离$fa$距离为$j-1$的点在第一次算子树距离为$j-2$的答案部分中已经算过一次，这里会重复计算。因此$d[i][j]$需要再减去$f[i][j-2]$。状态转移方程为$$d[i][j] = \sum d[fa][j-1] - f[i][j-2]$$在实际实现时，$d$和$f$数组可以用同一个，但是在第二次$DFS$进行容斥时需要倒序遍历（用到上一个时刻的状态量）。

#include<bits/stdc++.h>

const int N = 100010;

int n, k;
std::vector<std::vector<int>> g(N);
int c[N];
int f[N][25];

void dfs1(int u, int fa){
	for(auto v:g[u]){
		if(v == fa) continue;
		dfs1(v, u);
		for(int j = 1; j <= k; j ++ ){
			f[u][j] += f[v][j - 1];
		}
	}
}

void dfs2(int u, int fa){
	for(auto v:g[u]){
		if(v == fa) continue;
		for(int j = k; j >= 2; j -- ) f[v][j] -= f[v][j - 2];   //倒序遍历，类似于01背包
		for(int j = 1; j <= k; j ++ ) f[v][j] += f[u][j - 1];
		dfs2(v, u);
	}
}

int main(){
	std::ios::sync_with_stdio(false);
	std::cin.tie(0);

	std::cin >> n >> k;
	for(int i = 0; i < n - 1; i ++ ){
		int u, v;
		std::cin >> u >> v;
		g[u].push_back(v);
		g[v].push_back(u);
	}

	for(int i = 1; i <= n; i ++ ){
		std::cin >> c[i];
		f[i][0] = c[i];
	}

	dfs1(1, 0);
	dfs2(1, 0);

	for(int i = 1; i <= n; i ++ ){
		int ans = 0;
		for(int j = 0; j <= k; j ++ ){
			ans += f[i][j];
		}
		std::cout << ans << '\n';
	}

	return 0;
}