#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define int ll
#define endl '\n'
#define pb push_back
using pi = pair<int, int>;
template<class T> bool ckmin(T& a, const T& b) { return b < a ? a = b, 1 : 0; }
template<class T> bool ckmax(T& a, const T& b) { return a < b ? a = b, 1 : 0; }
/*
dp[i][j] - subarborele lui i daca mai am j de aruncat
*/
const int INF = 1e14;
const int N = 1e5 + 5, K = 105;
int n, k;
int a[N];
vector<int> adj[N];
vector<int> sum;
vector<vector<vector<int>>> max1, max2;
vector<vector<int>> dp;
void dfs(int u, int p) {
sum[u] = 0;
for (int v : adj[u]) {
if (v == p) continue;
dfs(v, u);
sum[u] += a[v];
}
for (int v : adj[u]) {
if (v == p) continue;
for (int i = 0; i <= k; ++i) {
// nu folosesc
if (dp[v][i] > max1[u][i][0]) {
swap(max1[u][i][0], max2[u][i][0]);
max1[u][i][0] = dp[v][i];
} else max2[u][i][0] = max(max2[u][i][0], dp[v][i]);
// folosesc
if (i == 0) continue;
if (dp[v][i - 1] > max1[u][i][1]) {
swap(max1[u][i][1], max2[u][i][1]);
max1[u][i][1] = dp[v][i - 1];
} else max2[u][i][1] = max(max2[u][i][1], dp[v][i - 1]);
}
}
for (int i = 0; i <= k; ++i) dp[u][i] = max({dp[u][i], max1[u][i][0], max1[u][i][1] + sum[u]});
}
int ans = 0;
void reroot(int u, int p) {
/*auto dpu = dp[u], dpp = dp[p];
auto sumu = sum[u], sump = sum[p];
auto max1u = max1[u], max2u = max2[u];*/
if (p != 0) {
// pentru p
sum[p] -= a[u];
for (int i = 0; i <= k; ++i) {
// nu folosesc
dp[p][i] = 0;
int maxi = max1[p][i][0];
if (dp[u][i] == maxi) maxi = max2[p][i][0];
dp[p][i] = max(dp[p][i], maxi);
// folosesc
if (i == 0) continue;
maxi = max1[p][i][1];
if (dp[u][i - 1] == maxi) maxi = max2[p][i][1];
dp[p][i] = max(dp[p][i], maxi + sum[p]);
}
// pentru u
sum[u] += a[p];
for (int i = 0; i <= k; ++i) {
// nu folosesc
dp[u][i] = 0;
if (dp[p][i] > max1[u][i][0]) {
swap(max1[u][i][0], max2[u][i][0]);
max1[u][i][0] = dp[p][i];
} else max2[u][i][0] = max(max2[u][i][0], dp[p][i]);
dp[u][i] = max(dp[u][i], max1[u][i][0]);
// folosesc
if (i == 0) continue;
if (dp[p][i - 1] > max1[u][i][1]) {
swap(max1[u][i][1], max2[u][i][1]);
max1[u][i][1] = dp[p][i - 1];
} else max2[u][i][1] = max(max2[u][i][1], dp[p][i - 1]);
dp[u][i] = max(dp[u][i], max1[u][i][1] + sum[u]);
}
}
ans = max(ans, dp[u][k]);
for (int v : adj[u]) {
if (v == p) continue;
reroot(v, u);
}
/*dp[u] = dpu;
dp[p] = dpp;
max1[u] = max1u;
max2[u] = max2u;
sum[u] = sumu;
sum[p] = sump;*.
}
int32_t main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL);
cin >> n >> k;
for (int i = 1; i <= n; ++i) cin >> a[i];
for (int i = 1; i < n; ++i) {
int u, v;
cin >> u >> v;
adj[u].pb(v);
adj[v].pb(u);
}
sum.assign(n + 1, 0);
max1.assign(n + 1, vector<vector<int>>(k + 1, vector<int>(2, -INF)));
max2.assign(n + 1, vector<vector<int>>(k + 1, vector<int>(2, -INF)));
dp.assign(n + 1, vector<int>(k + 1));
dfs(1, 0);
//for (int i = 0; i <= k; ++i) cout << max1[4][i][0] << " " << max2[4][i][0] << " | " << max1[4][i][1] << " " << max2[4][i][1] << endl;
reroot(1, 0);
cout << ans;
}
/*
good:
0 0 | -10000000000000000 -10000000000000000
16 0 | 0 0
27 0 | 16 0
bad:
0 0 | -10000000000000000 -10000000000000000
16 0 | 0 0
24 0 | 16 0
*/
Compilation message
chase.cpp:144:1: warning: "/*" within comment [-Wcomment]
144 | /*
|
chase.cpp: In function 'void reroot(ll, ll)':
chase.cpp:107:3: error: expected '}' at end of input
107 | }
| ^
chase.cpp:60:27: note: to match this '{'
60 | void reroot(int u, int p) {
| ^