题解 | 树上寻宝

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
#define ALL(a) a.begin(),a.end() 
const ll N = 1e5 + 7, M = 5;
struct p { int t, w; };
int n1[N];
bool vis[N];
vector<int> n2[N];
int main() {
	istream::sync_with_stdio(0), cin.tie(0);
	int n, k;cin >> n >> k;
	ll ans = 0;
	for (int i = 1;i <= n;++i) cin >> n1[i];
	for (int i = 0;i < n - 1;++i) {
		int u, v;cin >> u >> v;
		n2[u].push_back(v), n2[v].push_back(u);
	}
	queue<p> q;
	q.push({1, 0});
	vis[1] = 1;
	while (q.size()) {
		p tp = q.front();q.pop();
		if (tp.w > 2 * k) break;
		ans += n1[tp.t];
		for (int v : n2[tp.t]) if(!vis[v]) vis[v] = 1, q.push({v, tp.w + 1});
	}
	cout << ans;
	return 0;
}