[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

6 2 
5 1 
3 6 
2 4 
2 1 
3 2 
1 
2 
3 
4 
5 
6

样例输出 #1

15 
21 
16 
10 
8 
11

【数据范围】
对于 \(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;
}