[蓝桥杯 2025 省 AB/Python B 第二场] 树上寻宝

题目描述

有一棵包含 n 个结点的树，根结点是 1。 每个结点 i 上放着一个物品，价值为 \(w_i\)。

小蓝可以反复进行"寻宝"操作（次数不限）： 每次都从根结点 1 出发，最多走 k 步；每一步可以沿树上的边走 1 条边或连续走 2 条边； 走完后会自动拾取当前停留结点上的物品，然后被传送回根结点，开始下一次寻宝。

问：小蓝最终能获得的所有物品的总价值是多少？

关键理解： 一次寻宝最多能沿着根出发走 不超过 2k 条边的路径（因为每步可走 1 或 2 条边，共 k 步）。 因而所有与根结点距离 \(\le 2k\) 的结点，都可以在某一次作为终点被拿到其物品；寻宝次数不限，故这些结点的物品最终都会拿到。 于是答案就是：把所有满足 \(dist(1, i) \le 2k\) 的结点 i 的 \(w_i\) 之和。

输入格式

• 第一行：两个正整数 n, k。
• 第二行：n 个正整数 \(w_1, w_2, ..., w_n\)，表示各结点物品的价值。
• 接下来 \(n-1\) 行：每行两个正整数 \(u_i, v_i\)，表示结点 \(u_i\) 与 \(v_i\) 之间有一条边。

输出格式

输出一个整数，表示小蓝最终能获得的物品总价值。

输入输出样例 #1

输入 #1

8 2
6 3 3 1 5 4 3 4
1 2
2 3
2 4
4 5
5 6
6 7
7 8

输出 #1

22

说明/提示

样例说明

从根结点 1 出发，距离（按边数计）为：

• dist=0：1
• dist=1：2
• dist=2：3, 4
• dist=3：5
• dist=4：6
• dist=5：7
• dist=6：8

因为每次最多走 k=2 步、每步最多走 2 条边，所以一次最多走 2k=4 条边。 能作为终点的结点是距离 \( \le 4\) 的：1, 2, 3, 4, 5, 6。 这些结点的价值之和为 6+3+3+1+5+4 = 22。

评测用例规模与约定

• 对于 20% 的评测用例：\(1 \le n \le 15\)；
• 对于所有评测用例：\(0 \le k < n \le 10^5，1 \le w_i \le 10^6，1 \le u_i, v_i \le n\)。