先做子树内精确距离 DP，再做一次换根，把父亲方向的精确距离贡献传给儿子，最终累加 0..K 层即可。
OJ: luogu
题目 ID: P3047
难度:普及+/提高
标签:树形DP换根DP树动态规划
日期: 2026-06-21 03:32
题意
给一棵树，每个点有若干头牛。
对于每个点 i，要求统计：
- 距离
i不超过K的所有点上的牛数总和
思路
先看一个可以直接验证想法的朴素解：
cpp
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 30;
int n, k_limit;
vector<int> g[MAXN];
int cows[MAXN];
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
// 这是一个小数据精确暴力：
// 对每个点都做一次 BFS，统计距离不超过 K 的所有牛数。
cin >> n >> k_limit;
for (int i = 1; i <= n; i++) {
g[i].clear();
}
for (int i = 1; i < n; i++) {
int u, v;
cin >> u >> v;
g[u].push_back(v);
g[v].push_back(u);
}
for (int i = 1; i <= n; i++) {
cin >> cows[i];
}
for (int start = 1; start <= n; start++) {
vector<int> dist(n + 1, -1);
queue<int> q;
q.push(start);
dist[start] = 0;
while (!q.empty()) {
int u = q.front();
q.pop();
for (size_t i = 0; i < g[u].size(); i++) {
int v = g[u][i];
if (dist[v] != -1) {
continue;
}
dist[v] = dist[u] + 1;
q.push(v);
}
}
long long sum = 0;
for (int i = 1; i <= n; i++) {
if (dist[i] != -1 && dist[i] <= k_limit) {
sum += cows[i];
}
}
cout << sum << '\n';
}
return 0;
}brute.cpp 对每个点都做一次 BFS，统计距离不超过 K 的所有牛数。
这个方法完全正确，但无法处理大数据。
这题最关键的信息是：
K <= 20
所以我们可以直接做“按距离分层”的树形 DP。
先定义：
down_dp[u][d]：只看u子树时，与u距离恰好为d的牛数
这可以自底向上求：
down_dp[u][0] = cows[u]- 儿子
v的d-1层，会贡献给u的第d层
但仅靠子树信息还不够，因为答案还包含父亲方向、兄弟子树方向的牛。
所以再定义：
all_dp[u][d]：整棵树里，与u距离恰好为d的牛数
根节点直接有：
all_dp[1][d] = down_dp[1][d]
然后从父亲往儿子推：
对于儿子 v，距离 v 恰好为 d 的牛，分成两部分：
v子树内部的：down_dp[v][d]- 从
u方向过来的：all_dp[u][d-1] - down_dp[v][d-2]
第二项里减去 down_dp[v][d-2]，是为了去掉本来就在 v 子树里的那部分重复贡献。
最后，把 all_dp[u][0..K] 全加起来，就是点 u 的答案。
下面这棵样例树可以帮助理解“子树内”和“父亲方向”两类来源：
graph G {
5 -- 1;
3 -- 6;
2 -- 4;
2 -- 1;
3 -- 2;
}
比如对点 2 来说，距离不超过 2 的点既包括它自己子树里的 1,4，也包括往父亲和另一侧走到的 3,5,6。
这就是为什么只做子树 DP 不够，还需要第二遍换根传递。
DP 转移方程
核心状态：
down_dp[u][d] 与 all_dp[u][d]
核心转移：
all_dp[v][d]=down_dp[v][d]+all_dp[u][d-1]-down_dp[v][d-2]
答案收束：
ans[u]=sum_{d=0..K} all_dp[u][d]
代码
cpp
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100000 + 5;
const int MAXK = 20 + 5;
int n, k_limit;
vector<int> g[MAXN];
int cows[MAXN];
int parent_arr[MAXN];
long long down_dp[MAXN][MAXK];
long long all_dp[MAXN][MAXK];
long long answer_arr[MAXN];
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
cin >> n >> k_limit;
for (int i = 1; i <= n; i++) {
g[i].clear();
parent_arr[i] = 0;
answer_arr[i] = 0;
for (int d = 0; d <= k_limit; d++) {
down_dp[i][d] = 0;
all_dp[i][d] = 0;
}
}
for (int i = 1; i < n; i++) {
int u, v;
cin >> u >> v;
g[u].push_back(v);
g[v].push_back(u);
}
for (int i = 1; i <= n; i++) {
cin >> cows[i];
}
vector<int> order;
order.reserve(n);
stack<int> st;
st.push(1);
parent_arr[1] = 0;
while (!st.empty()) {
int u = st.top();
st.pop();
order.push_back(u);
for (size_t i = 0; i < g[u].size(); i++) {
int v = g[u][i];
if (v == parent_arr[u]) {
continue;
}
parent_arr[v] = u;
st.push(v);
}
}
// down_dp[u][d]：只看 u 子树时，和 u 距离恰好为 d 的牛数。
for (int idx = (int)order.size() - 1; idx >= 0; idx--) {
int u = order[idx];
down_dp[u][0] = cows[u];
for (size_t i = 0; i < g[u].size(); i++) {
int v = g[u][i];
if (v == parent_arr[u]) {
continue;
}
for (int d = 1; d <= k_limit; d++) {
down_dp[u][d] += down_dp[v][d - 1];
}
}
}
// all_dp[u][d]：整棵树里，和 u 距离恰好为 d 的牛数。
for (int d = 0; d <= k_limit; d++) {
all_dp[1][d] = down_dp[1][d];
}
for (size_t idx = 0; idx < order.size(); idx++) {
int u = order[idx];
for (size_t i = 0; i < g[u].size(); i++) {
int v = g[u][i];
if (v == parent_arr[u]) {
continue;
}
all_dp[v][0] = cows[v];
for (int d = 1; d <= k_limit; d++) {
// 先拿到“离 u 恰好 d-1”的所有牛，再减去来自 v 子树那部分，
// 剩下的就是通过父亲方向贡献给 v 的牛数。
long long from_parent_side = all_dp[u][d - 1];
if (d >= 2) {
from_parent_side -= down_dp[v][d - 2];
}
all_dp[v][d] = down_dp[v][d] + from_parent_side;
}
}
}
for (int u = 1; u <= n; u++) {
long long sum = 0;
for (int d = 0; d <= k_limit; d++) {
sum += all_dp[u][d];
}
answer_arr[u] = sum;
cout << answer_arr[u] << '\n';
}
return 0;
}复杂度
总共做两遍树上 DP，每次都要枚举 0..K 这一层距离。
所以时间复杂度是
总结
这题最值得记住的是：
- 当树上查询的“半径”很小的时候，可以直接把距离当成 DP 维度
然后通过：
- 一遍子树 DP
- 一遍换根 DP
把整棵树的信息补完整。
一图流解析
这张图把本题的建模、关键转移、实现检查和训练方法压缩到一页，适合读完正文后复盘。
